Add Gemini 2 Flash

google · Dec 11, 2024 · 33372d6 · 33372d6
1 parent b036b7a
commit 33372d6
Showing 1 changed file with 19 additions and 0 deletions.
diff --git a/llm_toolkit/models.py b/llm_toolkit/models.py
@@ -700,6 +700,25 @@ def chat_llm(self, client: ChatSession, prompt: prompts.Prompt) -> str:
     response = self._do_generate(client, prompt.get(), parameters_list) or ''
     return response
 
+class GeminiV1D5(GeminiModel):
+  """Gemini 1.5."""
+
+  _max_output_tokens = 8192
+  context_window = 2000000
+
+  name = 'vertex_ai_gemini-1-5'
+  _vertex_ai_model = 'gemini-1.5-pro-002'
+
+
+class GeminiV2Flash(GeminiModel):
+  """Gemini 2 flash."""
+
+  _max_output_tokens = 8192
+  context_window =  1048576
+
+  name = 'vertex_ai_gemini-2-flash'
+  _vertex_ai_model = 'gemini-2.0-flash-exp'
+
 
 class AIBinaryModel(GoogleModel):
   """A customized model hosted internally."""