feat(codegen): support model revision

Author: terryyz

bigcodebench/generate.py

@@ -139,6 +139,7 @@ def run_codegen(
     backend: str = "vllm",
     base_url: str = None,
     tp: int = 1,
+    revision: str = "main",
     trust_remote_code: bool = False,
     tokenizer_name: str = None,
     tokenizer_legacy: bool = False,
@@ -173,14 +174,15 @@ def run_codegen(
         response_prefix=response_prefix,
         base_url=base_url,
         tp=tp,
+        revision=revision,
         trust_remote_code=trust_remote_code,
         direct_completion=direct_completion,
         tokenizer_name=tokenizer_name,
         tokenizer_legacy=tokenizer_legacy
     )
 
     extra = "-" + subset if subset != "full" else ""
-    identifier = model.replace("/", "--") + f"--bigcodebench{extra}-{split}--{backend}-{temperature}-{n_samples}-sanitized_calibrated.jsonl"
+    identifier = model.replace("/", "--") + f"--{revision}--bigcodebench{extra}-{split}--{backend}-{temperature}-{n_samples}-sanitized_calibrated.jsonl"
 
     target_path = os.path.join(root, identifier)
 

bigcodebench/provider/__init__.py

@@ -12,6 +12,8 @@ def make_model(
     # instruction model only
     instruction_prefix: str = None,
     response_prefix: str = None,
+    # vllm and hf only
+    revision: str = "main",
     # vllm only
     tp: int = 1,
     direct_completion: bool = False,
@@ -32,6 +34,7 @@ def make_model(
             split=split,
             temperature=temperature,
             max_new_tokens=max_new_tokens,
+            revision=revision,
             dataset=dataset,
             direct_completion=direct_completion,
             tp=tp,
@@ -47,6 +50,7 @@ def make_model(
             split=split,
             temperature=temperature,
             max_new_tokens=max_new_tokens,
+            revision=revision,
             dataset=dataset,
             direct_completion=direct_completion,
             instruction_prefix=instruction_prefix,

bigcodebench/provider/base.py

@@ -12,6 +12,7 @@ def __init__(
         split: str,
         temperature: float = 0.8,
         max_new_tokens: int = 1280,
+        revision: str = "main",
         dtype: str = "bfloat16",  # default
         direct_completion: bool = False,
         trust_remote_code: bool = False,
@@ -29,6 +30,7 @@ def __init__(
         self.skip_special_tokens = False
         self.max_new_tokens = max_new_tokens
         self.dtype = dtype
+        self.revision = revision
         self.direct_completion = direct_completion
         self.trust_remote_code = trust_remote_code
         self.tokenizer_name = tokenizer_name

bigcodebench/provider/hf.py

@@ -27,6 +27,7 @@ def __init__(
             "trust_remote_code": self.trust_remote_code,
             "torch_dtype": getattr(torch, self.dtype),
             "attn_implementation": attn_implementation,  # "eager", "flash_attention_2", "sdpa"
+            "revision": self.revision,
         }
         self.skip_special_tokens = True
 

bigcodebench/provider/vllm.py

@@ -18,6 +18,7 @@ def __init__(self, name: str, dataset: str, tp: int, **kwargs) -> None:
             "tensor_parallel_size": int(os.getenv("VLLM_N_GPUS", tp)),
             "dtype": self.dtype,
             "trust_remote_code": self.trust_remote_code,
+            "revision": self.revision,
         }
         if self.tokenizer_name is None:
             self.tokenizer_name = self.name