feat: use google genai

Author: terryyz

Docker/Evaluate.Dockerfile

@@ -54,7 +54,7 @@ RUN pip install \
     rich \
     accelerate \
     anthropic \
-    google-generativeai \
+    google-genai \
     mistralai \
     openai \
     e2b

bigcodebench/gen/util/google_request.py

@@ -1,33 +1,47 @@
 import time
 
-import google.generativeai as genai
+from google import genai
 from google.api_core.exceptions import GoogleAPICallError, ResourceExhausted
 
 
 def make_request(
-    client: genai.GenerativeModel,
+    model: str,
+    client: genai.Client,
     message: str,
     temperature: float,
     n: int,
     max_new_tokens: int = 2048,
 ) -> genai.types.GenerateContentResponse:
     kwargs = {"temperature": temperature, "max_output_tokens": max_new_tokens}
 
-    if "-thinking-" in client.model_name:
+    if "-thinking-" in model:
         kwargs.pop("max_output_tokens")
-
-    response = client.generate_content(
-        [{"role": "user", "parts": [message]}],
-        generation_config=genai.types.GenerationConfig(
+    
+    response = client.models.generate_content(
+        model=model,
+        contents=message,
+        config=genai.types.GenerateContentConfig(
             candidate_count=n,
+            safety_settings=[
+                genai.types.SafetySetting(
+                    category='HARM_CATEGORY_DANGEROUS_CONTENT',
+                    threshold='BLOCK_NONE'
+                ),
+                genai.types.SafetySetting(
+                    category='HARM_CATEGORY_SEXUALLY_EXPLICIT',
+                    threshold='BLOCK_NONE'
+                ),
+                genai.types.SafetySetting(
+                    category='HARM_CATEGORY_HATE_SPEECH',
+                    threshold='BLOCK_NONE'
+                ),
+                genai.types.SafetySetting(
+                    category='HARM_CATEGORY_HARASSMENT',
+                    threshold='BLOCK_NONE'
+                ),
+            ],
             **kwargs
-        ),
-        safety_settings=[
-            {"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_NONE"},
-            {"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "BLOCK_NONE"},
-            {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_NONE"},
-            {"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_NONE"},
-        ],
+        ),            
     )
 
     return response

bigcodebench/provider/google.py

@@ -2,7 +2,7 @@
 from typing import List
 from tqdm import tqdm
 
-import google.generativeai as genai
+from google import genai
 
 from bigcodebench.provider.base import DecoderBase
 from bigcodebench.gen.util.google_request import make_auto_request
@@ -12,8 +12,8 @@
 class GoogleDecoder(DecoderBase):
     def __init__(self, name: str, **kwargs):
         super().__init__(name, **kwargs)
-        genai.configure(api_key=os.getenv("GOOGLE_API_KEY"))
-        self.client = genai.GenerativeModel(name)
+        self.model = name
+        self.client = genai.Client(api_key=os.getenv("GOOGLE_API_KEY"))
 
     def codegen(
         self, prompts: List[str], do_sample: bool = True, num_samples: int = 200
@@ -34,7 +34,8 @@ def codegen(
                 tokenizer=None,
             )
             ret = make_auto_request(
-                self.client,
+                model=self.model,
+                client=self.client,
                 message=message,
                 n=num_samples,
                 temperature=self.temperature,

setup.cfg

@@ -35,7 +35,7 @@ install_requires =
     rich
     accelerate>=0.30.1
     anthropic>=0.26.1
-    google-generativeai>=0.5.4
+    google-genai
     mistralai>=0.2.0,<1.0.0
     openai>=1.11.1
     e2b