fix(generate): use 2nd part for gemini thinking models

bigcode-project · Dec 20, 2024 · fd3cbc8 · fd3cbc8
1 parent ba18aaf
commit fd3cbc8
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/bigcodebench/provider/google.py b/bigcodebench/provider/google.py
@@ -43,7 +43,10 @@ def codegen(
             for candidate in ret.candidates:
                 parts = candidate.content.parts
                 if parts:
-                    outputs.append(parts[0].text)
+                    if "-thinking-" in self.name:
+                        outputs.append(parts[1].text)
+                    else:
+                        outputs.append(parts[0].text)
                 else:
                     print("Empty response!")
                     outputs.append("")