Merge pull request #37 from allegro/better-gemini-model-names-support

Better gemini model names support

allms/models/vertexai_gemini.py

@@ -1,17 +1,20 @@
 import typing
 from asyncio import AbstractEventLoop
+from typing import Optional
 
 from langchain_core.prompts import ChatPromptTemplate
 from langchain_google_vertexai import VertexAI
 from vertexai.preview import tokenization
-from typing import Optional
+from vertexai.tokenization._tokenizers import Tokenizer
 
 from allms.defaults.general_defaults import GeneralDefaults
 from allms.defaults.vertex_ai import GeminiModelDefaults
 from allms.domain.configuration import VertexAIConfiguration
 from allms.domain.input_data import InputData
-from allms.models.vertexai_base import CustomVertexAI
 from allms.models.abstract import AbstractModel
+from allms.models.vertexai_base import CustomVertexAI
+
+BASE_GEMINI_MODEL_NAMES = ["gemini-1.0-pro", "gemini-1.5-pro", "gemini-1.5-flash"]
 
 
 class VertexAIGeminiModel(AbstractModel):
@@ -33,7 +36,7 @@ def __init__(
         self._verbose = verbose
         self._config = config
 
-        self._gcp_tokenizer = tokenization.get_tokenizer_for_model(self._config.gemini_model_name)
+        self._gcp_tokenizer = self._get_gcp_tokenizer(self._config.gemini_model_name)
 
         super().__init__(
             temperature=temperature,
@@ -67,3 +70,13 @@ def _get_model_response_tokens_number(self, model_response: typing.Optional[str]
             return self._gcp_tokenizer.count_tokens(model_response).total_tokens
         return 0
 
+    @staticmethod
+    def _get_gcp_tokenizer(model_name) -> Tokenizer:
+        try:
+            return tokenization.get_tokenizer_for_model(model_name)
+        except ValueError:
+            for base_model_name in BASE_GEMINI_MODEL_NAMES:
+                if model_name.startswith(base_model_name):
+                    return tokenization.get_tokenizer_for_model(base_model_name)
+            raise
+

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "allms"
-version = "1.0.9"
+version = "1.0.10"
 description = ""
 authors = ["Allegro Opensource <[email protected]>"]
 readme = "README.md"

tests/test_end_to_end.py

@@ -1,7 +1,8 @@
 import re
 
-from unittest.mock import patch
-from langchain.prompts import ChatPromptTemplate, HumanMessagePromptTemplate, PromptTemplate, SystemMessagePromptTemplate
+import pytest
+from langchain.prompts import ChatPromptTemplate, HumanMessagePromptTemplate, PromptTemplate, \
+    SystemMessagePromptTemplate
 
 from allms.constants.input_data import IODataConstants
 from allms.domain.configuration import VertexAIConfiguration
@@ -171,6 +172,40 @@ def test_gemini_specific_args_are_passed_to_model(self):
         assert gemini_model._llm.model_name == gemini_model_name
         assert gemini_model._llm.safety_settings == gemini_safety_settings
 
+    @pytest.mark.parametrize(
+        "model_name", [
+            "gemini-1.0-pro", "gemini-1.5-pro", "gemini-1.5-flash","gemini-1.0-pro-001", "gemini-1.0-pro-002",
+            "gemini-1.5-pro-001", "gemini-1.5-flash-001", "gemini-1.5-pro-preview-0514"
+        ]
+    )
+    def test_correct_gemini_model_name_work(self, model_name):
+        # GIVEN
+        model_config = VertexAIConfiguration(
+            cloud_project="dummy-project-id",
+            cloud_location="us-central1",
+            gemini_model_name=model_name,
+        )
+
+        # WHEN & THEN
+        VertexAIGeminiModel(config=model_config)
+
+    @pytest.mark.parametrize(
+        "model_name", [
+            "gemini-2.0-pro", "geminis-1.5-pro", "gemini-flash", "gemini-1.5-preview-pro", "gpt4"
+        ]
+    )
+    def test_incorrect_gemini_model_name_fail(self, model_name):
+        # GIVEN
+        model_config = VertexAIConfiguration(
+            cloud_project="dummy-project-id",
+            cloud_location="us-central1",
+            gemini_model_name=model_name,
+        )
+
+        # WHEN & THEN
+        with pytest.raises(ValueError, match=f"Model {model_name} is not supported."):
+            VertexAIGeminiModel(config=model_config)
+
     def test_model_times_out(
             self,
             mock_aioresponse,