[Model] Support GeminiFlash2-0

open-compass · Dec 13, 2024 · 2d1c026 · 2d1c026
1 parent 27d8ca8
commit 2d1c026
Show file tree

Hide file tree

Showing 2 changed files with 1 addition and 3 deletions.
diff --git a/vlmeval/api/gemini.py b/vlmeval/api/gemini.py
@@ -22,8 +22,6 @@ def __init__(self,
                  project_id='vlmeval',
                  **kwargs):
 
-        assert model in ['gemini-1.0-pro', 'gemini-1.5-pro', 'gemini-1.5-flash']
-
         self.model = model
         self.fail_msg = 'Failed to obtain answer via API. '
         self.max_tokens = max_tokens
@@ -77,7 +75,6 @@ def generate_inner(self, inputs, **kwargs) -> str:
             if pure_text and self.model == 'gemini-1.0-pro':
                 model = genai.GenerativeModel('gemini-1.0-pro')
             else:
-                assert self.model in ['gemini-1.5-pro', 'gemini-1.5-flash']
                 model = genai.GenerativeModel(self.model)
 
             messages = self.build_msgs_genai(inputs)

diff --git a/vlmeval/config.py b/vlmeval/config.py
@@ -65,6 +65,7 @@
     'GeminiPro1-0': partial(GeminiProVision, model='gemini-1.0-pro', temperature=0, retry=10),  # now GeminiPro1-0 is only supported by vertex backend
     'GeminiPro1-5': partial(GeminiProVision, model='gemini-1.5-pro', temperature=0, retry=10),
     'GeminiFlash1-5': partial(GeminiProVision, model='gemini-1.5-flash', temperature=0, retry=10),
+    'GeminiFlash2-0': partial(GeminiProVision, model='gemini-2.0-flash-exp', temperature=0, retry=10),
     'GeminiPro1-5-002': partial(GPT4V, model='gemini-1.5-pro-002', temperature=0, retry=10),  # Internal Use Only
     'GeminiFlash1-5-002': partial(GPT4V, model='gemini-1.5-flash-002', temperature=0, retry=10),  # Internal Use Only
     # Qwen-VL