From 0748f2e8d7a1fb7298af56142d58f70d33773e79 Mon Sep 17 00:00:00 2001
From: Huu Le <39040748+leehuwuj@users.noreply.github.com>
Date: Fri, 14 Jun 2024 14:18:23 +0700
Subject: [PATCH] remove gemini model map (#128)
---
.changeset/silver-masks-speak.md | 5 +++++
.../types/streaming/fastapi/app/settings.py | 18 ++++--------------
2 files changed, 9 insertions(+), 14 deletions(-)
create mode 100644 .changeset/silver-masks-speak.md
diff --git a/.changeset/silver-masks-speak.md b/.changeset/silver-masks-speak.md
new file mode 100644
index 00000000..704befc9
--- /dev/null
+++ b/.changeset/silver-masks-speak.md
@@ -0,0 +1,5 @@
+---
+"create-llama": patch
+---
+
+Remove hard-coded Gemini supported models
diff --git a/templates/types/streaming/fastapi/app/settings.py b/templates/types/streaming/fastapi/app/settings.py
index 3f138b10..87c591af 100644
--- a/templates/types/streaming/fastapi/app/settings.py
+++ b/templates/types/streaming/fastapi/app/settings.py
@@ -115,18 +115,8 @@ def init_gemini():
from llama_index.llms.gemini import Gemini
from llama_index.embeddings.gemini import GeminiEmbedding
- model_map: Dict[str, str] = {
- "gemini-1.5-pro-latest": "models/gemini-1.5-pro-latest",
- "gemini-pro": "models/gemini-pro",
- "gemini-pro-vision": "models/gemini-pro-vision",
- }
+ model_name = f"models/{os.getenv('MODEL')}"
+ embed_model_name = f"models/{os.getenv('EMBEDDING_MODEL')}"
- embed_model_map: Dict[str, str] = {
- "embedding-001": "models/embedding-001",
- "text-embedding-004": "models/text-embedding-004",
- }
-
- Settings.llm = Gemini(model=model_map[os.getenv("MODEL")])
- Settings.embed_model = GeminiEmbedding(
- model_name=embed_model_map[os.getenv("EMBEDDING_MODEL")]
- )
+ Settings.llm = Gemini(model=model_name)
+ Settings.embed_model = GeminiEmbedding(model_name=embed_model_name)
--
GitLab