requirement: google-cloud-aiplatform = ">=1.57.0" modified

allegro · Aug 8, 2024 · a8d8b3f · a8d8b3f
1 parent 7f6428f
commit a8d8b3f
Show file tree

Hide file tree

Showing 4 changed files with 14 additions and 15 deletions.
diff --git a/allms/defaults/vertex_ai.py b/allms/defaults/vertex_ai.py
@@ -10,7 +10,7 @@ class PalmModelDefaults:
 
 
 class GeminiModelDefaults:
-    GCP_MODEL_NAME = "gemini-1.0-pro-001"
+    GCP_MODEL_NAME = "gemini-1.5-flash-001"
     MODEL_TOTAL_MAX_TOKENS = 30720
     MAX_OUTPUT_TOKENS = 2048
     TEMPERATURE = 0.0

diff --git a/poetry.lock b/poetry.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -9,7 +9,7 @@ packages = [{include = "allms"}]
 [tool.poetry.dependencies]
 python = ">=3.8.1,<4.0"
 fsspec = "^2023.6.0"
-google-cloud-aiplatform = "^1.47.0"
+google-cloud-aiplatform = ">=1.57.0"
 pydash = "^7.0.6"
 transformers = "^4.34.1"
 pydantic = "1.10.13"

diff --git a/tests/test_end_to_end.py b/tests/test_end_to_end.py
@@ -72,7 +72,6 @@ def test_model_is_queried_successfully(
         assert list(map(lambda output: int(output[IODataConstants.GENERATED_TOKENS_NUMBER]), expected_output)) == list(
             map(lambda example: example.number_of_generated_tokens, parsed_responses))
 
-
     def test_prompt_is_not_modified_for_open_source_models(self, mock_aioresponse, models, mocker):
         # GIVEN
         open_source_models = ["azure_llama2", "azure_mistral", "vertex_gemma"]
@@ -150,7 +149,7 @@ def test_prompt_is_not_modified_for_open_source_models(self, mock_aioresponse, m
 
     def test_gemini_specific_args_are_passed_to_model(self):
         # GIVEN
-        gemini_model_name = "gemini-1.0-pro-001"
+        gemini_model_name = "gemini-1.5-pro-001"
         gemini_safety_settings = {
             HarmCategory.HARM_CATEGORY_UNSPECIFIED: HarmBlockThreshold.BLOCK_NONE,
             HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,