factor out VertaxAIGemini1_5 base class. It allows user to pass in an…

…y Gemini 1.5 model name. PiperOrigin-RevId: 682428314
google · Oct 4, 2024 · ec3cfbe · ec3cfbe
1 parent 0a2a548
commit ec3cfbe
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 43 deletions.
diff --git a/langfun/core/llms/__init__.py b/langfun/core/llms/__init__.py
@@ -102,6 +102,7 @@
 from langfun.core.llms.groq import GroqGemma7B_IT
 
 from langfun.core.llms.vertexai import VertexAI
+from langfun.core.llms.vertexai import VertexAIGemini1_5
 from langfun.core.llms.vertexai import VertexAIGeminiPro1_5
 from langfun.core.llms.vertexai import VertexAIGeminiPro1_5_Latest
 from langfun.core.llms.vertexai import VertexAIGeminiPro1_5_001

diff --git a/langfun/core/llms/vertexai.py b/langfun/core/llms/vertexai.py
@@ -422,101 +422,78 @@ def get_text_generation_model(
 ]
 
 
-class VertexAIGeminiPro1_5_Latest(VertexAI):  # pylint: disable=invalid-name
-  """Vertex AI Gemini 1.5 Pro model."""
+class VertexAIGemini1_5(VertexAI):  # pylint: disable=invalid-name
+  """Vertex AI Gemini 1.5 model."""
 
-  model = 'gemini-1.5-pro-latest'
   supported_modalities = (
       _DOCUMENT_TYPES + _IMAGE_TYPES + _AUDIO_TYPES + _VIDEO_TYPES
   )
 
 
-class VertexAIGeminiPro1_5(VertexAI):  # pylint: disable=invalid-name
+class VertexAIGeminiPro1_5_Latest(VertexAIGemini1_5):  # pylint: disable=invalid-name
+  """Vertex AI Gemini 1.5 Pro model."""
+
+  model = 'gemini-1.5-pro-latest'
+
+
+class VertexAIGeminiPro1_5(VertexAIGemini1_5):  # pylint: disable=invalid-name
   """Vertex AI Gemini 1.5 Pro model."""
 
   model = 'gemini-1.5-pro'
-  supported_modalities = (
-      _DOCUMENT_TYPES + _IMAGE_TYPES + _AUDIO_TYPES + _VIDEO_TYPES
-  )
 
 
-class VertexAIGeminiPro1_5_002(VertexAI):  # pylint: disable=invalid-name
+class VertexAIGeminiPro1_5_002(VertexAIGemini1_5):  # pylint: disable=invalid-name
   """Vertex AI Gemini 1.5 Pro model."""
 
   model = 'gemini-1.5-pro-002'
-  supported_modalities = (
-      _DOCUMENT_TYPES + _IMAGE_TYPES + _AUDIO_TYPES + _VIDEO_TYPES
-  )
 
 
-class VertexAIGeminiPro1_5_001(VertexAI):  # pylint: disable=invalid-name
+class VertexAIGeminiPro1_5_001(VertexAIGemini1_5):  # pylint: disable=invalid-name
   """Vertex AI Gemini 1.5 Pro model."""
 
   model = 'gemini-1.5-pro-001'
-  supported_modalities = (
-      _DOCUMENT_TYPES + _IMAGE_TYPES + _AUDIO_TYPES + _VIDEO_TYPES
-  )
 
 
-class VertexAIGeminiPro1_5_0514(VertexAI):  # pylint: disable=invalid-name
+class VertexAIGeminiPro1_5_0514(VertexAIGemini1_5):  # pylint: disable=invalid-name
   """Vertex AI Gemini 1.5 Pro preview model."""
 
   model = 'gemini-1.5-pro-preview-0514'
-  supported_modalities = (
-      _DOCUMENT_TYPES + _IMAGE_TYPES + _AUDIO_TYPES + _VIDEO_TYPES
-  )
 
 
-class VertexAIGeminiPro1_5_0409(VertexAI):  # pylint: disable=invalid-name
+class VertexAIGeminiPro1_5_0409(VertexAIGemini1_5):  # pylint: disable=invalid-name
   """Vertex AI Gemini 1.5 Pro preview model."""
 
   model = 'gemini-1.5-pro-preview-0409'
-  supported_modalities = (
-      _DOCUMENT_TYPES + _IMAGE_TYPES + _AUDIO_TYPES + _VIDEO_TYPES
-  )
 
 
-class VertexAIGeminiFlash1_5_Latest(VertexAI):  # pylint: disable=invalid-name
+class VertexAIGeminiFlash1_5_Latest(VertexAIGemini1_5):  # pylint: disable=invalid-name
   """Vertex AI Gemini 1.5 Flash model."""
 
   model = 'gemini-1.5-flash-latest'
-  supported_modalities = (
-      _DOCUMENT_TYPES + _IMAGE_TYPES + _AUDIO_TYPES + _VIDEO_TYPES
-  )
 
 
-class VertexAIGeminiFlash1_5(VertexAI):  # pylint: disable=invalid-name
+class VertexAIGeminiFlash1_5(VertexAIGemini1_5):  # pylint: disable=invalid-name
   """Vertex AI Gemini 1.5 Flash model."""
+
   model = 'gemini-1.5-flash'
-  supported_modalities = (
-      _DOCUMENT_TYPES + _IMAGE_TYPES + _AUDIO_TYPES + _VIDEO_TYPES
-  )
 
 
-class VertexAIGeminiFlash1_5_002(VertexAI):  # pylint: disable=invalid-name
+class VertexAIGeminiFlash1_5_002(VertexAIGemini1_5):  # pylint: disable=invalid-name
   """Vertex AI Gemini 1.5 Flash model."""
 
   model = 'gemini-1.5-flash-002'
-  supported_modalities = (
-      _DOCUMENT_TYPES + _IMAGE_TYPES + _AUDIO_TYPES + _VIDEO_TYPES
-  )
 
 
-class VertexAIGeminiFlash1_5_001(VertexAI):  # pylint: disable=invalid-name
+class VertexAIGeminiFlash1_5_001(VertexAIGemini1_5):  # pylint: disable=invalid-name
   """Vertex AI Gemini 1.5 Flash model."""
+
   model = 'gemini-1.5-flash-001'
-  supported_modalities = (
-      _DOCUMENT_TYPES + _IMAGE_TYPES + _AUDIO_TYPES + _VIDEO_TYPES
-  )
 
 
-class VertexAIGeminiFlash1_5_0514(VertexAI):  # pylint: disable=invalid-name
+class VertexAIGeminiFlash1_5_0514(VertexAIGemini1_5):  # pylint: disable=invalid-name
   """Vertex AI Gemini 1.5 Flash preview model."""
 
   model = 'gemini-1.5-flash-preview-0514'
-  supported_modalities = (
-      _DOCUMENT_TYPES + _IMAGE_TYPES + _AUDIO_TYPES + _VIDEO_TYPES
-  )
 
 
 class VertexAIGeminiPro1(VertexAI):  # pylint: disable=invalid-name