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)