Update to Gemini 2.0

Author: tomasruizt

llmlib/llmlib/gemini/media_description.py llmlib/llmlib/gemini/gemini_code.pyllmlib/llmlib/gemini/media_description.py renamed to llmlib/llmlib/gemini/gemini_code.py

@@ -8,7 +8,6 @@
 from logging import getLogger
 from pathlib import Path
 import tempfile
-from typing import Literal
 from google.cloud import storage
 from google.cloud.storage import transfer_manager
 import proto
@@ -19,6 +18,7 @@
     HarmBlockThreshold,
     GenerationResponse,
 )
+from enum import StrEnum
 from ..base_llm import LLM, Message
 from ..error_handling import notify_bugsnag
 
@@ -27,7 +27,7 @@
 logger = getLogger(__name__)
 
 project_id = "css-lehrbereich"  # from google cloud console
-frankfurt = "europe-west3"  # https://cloud.google.com/about/locations#europe
+location = "europe-west1"  # https://cloud.google.com/about/locations#europe
 
 
 class Buckets:
@@ -40,18 +40,23 @@ def storage_uri(bucket: str, blob_name: str) -> str:
     return "gs://%s/%s" % (bucket, blob_name)
 
 
-class Models:
-    gemini_pro = "models/gemini-1.5-pro"
-    gemini_flash = "models/gemini-1.5-flash"
+class GeminiModels(StrEnum):
+    gemini_15_pro = "models/gemini-1.5-pro"
+    gemini_20_flash = "models/gemini-2.0-flash"
+    gemini_20_flash_lite = "models/gemini-2.0-flash-lite"
 
 
-available_models = [Models.gemini_pro, Models.gemini_flash]
+available_models = [
+    GeminiModels.gemini_15_pro,
+    GeminiModels.gemini_20_flash,
+    GeminiModels.gemini_20_flash_lite,
+]
 
 
 @dataclass
 class Request:
     media_files: list[Path]
-    model_name: Literal[Models.gemini_pro, Models.gemini_flash] = Models.gemini_pro
+    model_name: GeminiModels = GeminiModels.gemini_15_pro
     prompt: str = "Describe this video in detail."
     max_output_tokens: int = 1000
 
@@ -101,7 +106,7 @@ def fetch_media_description(req: Request) -> str:
 
 
 def init_vertex() -> None:
-    vertexai.init(project=project_id, location=frankfurt)
+    vertexai.init(project=project_id, location=location)
 
 
 def mime_type(file_name: str) -> str:
@@ -174,11 +179,11 @@ class ResponseRefusedException(Exception):
 
 @dataclass
 class GeminiAPI(LLM):
-    model_id: str = Models.gemini_pro
+    model_id: str = GeminiModels.gemini_20_flash_lite
     max_output_tokens: int = 1000
 
     requires_gpu_exclusively = False
-    model_ids = [Models.gemini_pro, Models.gemini_flash]
+    model_ids = available_models
 
     def complete_msgs(self, msgs: list[Message]) -> str:
         if len(msgs) != 1:

llmlib/llmlib/runtime.py

@@ -1,6 +1,6 @@
 from .replicate_api import Apollo7B
 from .internvl import InternVL
-from .gemini.media_description import GeminiAPI
+from .gemini.gemini_code import GeminiAPI
 from .gemma import PaliGemma2
 from .minicpm import MiniCPM
 from .llama3 import LLama3Vision8B

tests/test_gemini.py

@@ -1,5 +1,5 @@
 from pathlib import Path
-from llmlib.gemini.media_description import GeminiAPI, Models, Request
+from llmlib.gemini.gemini_code import GeminiAPI, GeminiModels, Request
 import pytest
 
 from tests.helpers import (
@@ -23,7 +23,7 @@ def test_gemini_vision():
         assert path.exists()
 
     req = Request(
-        model_name=Models.gemini_flash,
+        model_name=GeminiModels.gemini_20_flash,
         media_files=files,
         prompt="Describe this combined images/audio/text in detail.",
     )
@@ -35,7 +35,7 @@ def test_gemini_vision():
 
 @pytest.mark.skipif(condition=is_ci(), reason="Avoid costs")
 def test_gemini_vision_using_interface():
-    model = GeminiAPI(model_id=Models.gemini_flash, max_output_tokens=50)
+    model = GeminiAPI(model_id=GeminiModels.gemini_20_flash_lite, max_output_tokens=50)
     assert_model_knows_capital_of_france(model)
     assert_model_recognizes_pyramid_in_image(model)
     assert_model_recognizes_afd_in_video(model)