Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

factor out VertaxAIGemini1_5 base class. It allows user to pass in any Gemini 1.5 model name. #306

Merged
merged 1 commit into from
Oct 4, 2024
Merged

factor out VertaxAIGemini1_5 base class. It allows user to pass in any Gemini 1.5 model name. #306

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions langfun/core/llms/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
63 changes: 20 additions & 43 deletions langfun/core/llms/vertexai.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Loading