feat: support anthropic extended thinking

Author: terryyz

bigcodebench/gen/util/anthropic_request.py

@@ -16,7 +16,18 @@ def make_auto_request(client: anthropic.Client, *args, **kwargs) -> Message:
         try:
             signal.signal(signal.SIGALRM, handler)
             signal.alarm(100)
-            ret = client.messages.create(*args, **kwargs)
+            if "reasoning_budget" in kwargs and "reasoning_beta" in kwargs:
+                ret = client.beta.messages.create(
+                    *args, 
+                    **kwargs, 
+                    thinking = {
+                        "type": "enabled",
+                        "budget": kwargs["reasoning_budget"],
+                    },
+                    betas=[kwargs["reasoning_beta"]]
+                )
+            else:
+                ret = client.messages.create(*args, **kwargs)
             signal.alarm(0)
         except anthropic.RateLimitError:
             print("Rate limit exceeded. Waiting...")

bigcodebench/generate.py

@@ -132,7 +132,11 @@ def run_codegen(
     temperature: float = 0.0,
     max_new_tokens: int = 1280,
     greedy: bool = False,
+    # openai
     reasoning_effort: str = "medium",
+    # anthropic
+    reasoning_budget: int = 0,
+    reasoning_beta: str = "output-128k-2025-02-19",
     strip_newlines: bool = False,
     direct_completion: bool = False,
     resume: bool = True,
@@ -173,6 +177,8 @@ def run_codegen(
         temperature=temperature,
         max_new_tokens=max_new_tokens,
         reasoning_effort=reasoning_effort,
+        reasoning_budget=reasoning_budget,
+        reasoning_beta=reasoning_beta,
         instruction_prefix=instruction_prefix,
         response_prefix=response_prefix,
         prefill=not skip_prefill,
@@ -186,8 +192,11 @@ def run_codegen(
     )
 
     extra = "-" + subset if subset != "full" else ""
-    if reasoning_effort and model.startswith("o1-") or model.startswith("o3-") or model.endswith("-reasoner"):
+    if backend == "openai" and reasoning_effort and model.startswith("o1-") or model.startswith("o3-") or model.endswith("-reasoner"):
         model = model + f"--{reasoning_effort}"
+    
+    if backend == "anthropic" and reasoning_budget and reasoning_beta:
+        model = model + f"--{reasoning_budget}-{reasoning_beta}"
 
     if skip_prefill:
         identifier = model.replace("/", "--") + "--skip_prefill" + f"--{revision}--bigcodebench{extra}-{split}--{backend}-{temperature}-{n_samples}-sanitized_calibrated.jsonl"

bigcodebench/provider/__init__.py

@@ -9,8 +9,11 @@ def make_model(
     dataset: str = "bigcodebench",
     temperature: float = 0.0,
     max_new_tokens: int = 1280,
-    # o1 and o3 only
+    # openai only
     reasoning_effort: str = "medium",
+    # anthropic only
+    reasoning_budget: int = 0,
+    reasoning_beta: str = "output-128k-2025-02-19",
     # instruction model only
     instruction_prefix: str = None,
     response_prefix: str = None,
@@ -118,6 +121,8 @@ def make_model(
             split=split,
             temperature=temperature,
             max_new_tokens=max_new_tokens,
+            reasoning_budget=reasoning_budget,
+            reasoning_beta=reasoning_beta,
             instruction_prefix=instruction_prefix,
             response_prefix=response_prefix,
         )

bigcodebench/provider/anthropic.py

@@ -9,9 +9,11 @@
 from bigcodebench.provider.utility import make_raw_chat_prompt
 
 class AnthropicDecoder(DecoderBase):
-    def __init__(self, name: str, **kwargs) -> None:
+    def __init__(self, name: str, reasoning_budget: int = 0, reasoning_beta: str = "output-128k-2025-02-19", **kwargs) -> None:
         super().__init__(name, **kwargs)
         self.client = anthropic.Anthropic(api_key=os.getenv("ANTHROPIC_KEY"))
+        self.reasoning_budget = reasoning_budget
+        self.reasoning_beta = reasoning_beta
 
     def codegen(
         self, prompts: List[str], do_sample: bool = True, num_samples: int = 200
@@ -43,6 +45,8 @@ def codegen(
                     max_tokens=self.max_new_tokens,
                     temperature=self.temperature,
                     stop_sequences=self.eos,
+                    reasoning_budget=self.reasoning_budget,
+                    reasoning_beta=self.reasoning_beta,
                 )
                 outputs.append(ret.content[0].text)
             all_outputs.append(outputs)