fix: make google api do n samples

Author: terryyz

bigcodebench/gen/util/google_request.py

@@ -5,13 +5,17 @@
 
 
 def make_request(
-    client: genai.GenerativeModel, temperature, messages, max_new_tokens=2048
+    client: genai.GenerativeModel,
+    messages: List,
+    temperature: float,
+    n: int,
+    max_new_tokens: int = 2048,
 ) -> genai.types.GenerateContentResponse:
     messages = [{"role": m["role"], "parts": [m["content"]]} for m in messages]
     response = client.generate_content(
         messages,
         generation_config=genai.types.GenerationConfig(
-            candidate_count=1,
+            candidate_count=n,
             max_output_tokens=max_new_tokens,
             temperature=temperature,
         ),

bigcodebench/provider/google.py

@@ -36,10 +36,9 @@ def codegen(
             ret = make_auto_request(
                 self.client,
                 message,
-                self.name,
                 n=num_samples,
-                max_tokens=self.max_new_tokens,
                 temperature=self.temperature,
+                max_new_tokens=self.max_new_tokens,
             )
             for candidate in ret.candidates:
                 parts = candidate.content.parts