feat: add customized tokenizer

Author: terryyz

bigcodebench/evaluate.py

@@ -281,6 +281,32 @@ def stucking_checker():
     if not os.path.isfile(result_path):
         with open(result_path, "w") as f:
             json.dump(results, f, indent=2)
+    
+    pass_at_k_path = result_path.replace("_eval_results.json", "_pass_at_k.json")
+    pass_at_k["model"] = flags.samples.split("/")[-1].replace(".jsonl", "")
+    pass_at_k["subset"] = flags.subset
+
+    def save_pass_at_k():
+        with open(pass_at_k_path, "w") as f:
+            json.dump(pass_at_k, f, indent=2)
+
+    if os.path.isfile(pass_at_k_path):
+        saved_pass_at_k = json.load(open(pass_at_k_path, "r"))
+        # compare saved_pass_at_k with pass_at_k
+        for k in saved_pass_at_k.keys():
+            if pass_at_k[k] != saved_pass_at_k[k]:
+                cprint(f"Warning: {k} is different from the saved one", "yellow")
+                
+        # ask user whether to save the pass@k
+        decision = ""
+        while decision.lower() not in ["y", "n"]:
+            print(f"Save pass@k to {pass_at_k_path}? [Y/N]")
+            decision = input()
+        if decision.lower() == "y":
+            save_pass_at_k()
+            
+    else:
+        save_pass_at_k()
 
 
 def main():

bigcodebench/generate.py

@@ -119,6 +119,7 @@ def main():
     parser.add_argument("--base_url", default=None, type=str)
     parser.add_argument("--tp", default=1, type=int)
     parser.add_argument("--trust_remote_code", action="store_true")
+    parser.add_argument("--tokenizer_name", default=None, type=str)
     args = parser.parse_args()
 
     if args.greedy and (args.temperature != 0 or args.bs != 1 or args.n_samples != 1)\
@@ -142,7 +143,8 @@ def main():
         temperature=args.temperature,
         base_url=args.base_url,
         tp=args.tp,
-        trust_remote_code=args.trust_remote_code
+        trust_remote_code=args.trust_remote_code,
+        tokenizer_name=args.tokenizer_name
     )
 
     extra = "-" + args.subset if args.subset != "full" else ""

bigcodebench/model.py

@@ -91,6 +91,7 @@ def __init__(
         max_new_tokens: int = 1280,
         dtype: str = "bfloat16",  # default
         trust_remote_code: bool = False,
+        tokenizer_name: str = None,
     ) -> None:
         print("Initializing a decoder model: {} ...".format(name))
         self.name = name
@@ -101,6 +102,7 @@ def __init__(
         self.max_new_tokens = max_new_tokens
         self.dtype = dtype
         self.trust_remote_code = trust_remote_code
+        self.tokenizer_name = tokenizer_name
 
     @abstractmethod
     def codegen(
@@ -128,8 +130,10 @@ def __init__(self, name: str, dataset: str, tp: int, **kwargs) -> None:
             "dtype": self.dtype,
             "trust_remote_code": self.trust_remote_code,
         }
-
-        self.tokenizer = AutoTokenizer.from_pretrained(self.name, legacy=False, **kwargs)
+        if self.tokenizer_name is None:
+            self.tokenizer_name = self.name
+        
+        self.tokenizer = AutoTokenizer.from_pretrained(self.tokenizer_name, legacy=False, **kwargs)
         if self.tokenizer.chat_template is None:
             self.eos += extra_eos_for_direct_completion(dataset)
         self.llm = LLM(model=name, max_model_len=2048, **kwargs)
@@ -185,9 +189,11 @@ def __init__(self, name: str, dataset: str, **kwargs):
         kwargs["torch_dtype"] = getattr(torch, self.dtype)
         self.skip_special_tokens = True
 
-        print(f"{kwargs = }")
-
-        self.tokenizer = AutoTokenizer.from_pretrained(name, legacy=False, **kwargs)
+        print(f"{kwargs = }", self.tokenizer_name)
+        if self.tokenizer_name is None:
+            self.tokenizer_name = self.name
+        
+        self.tokenizer = AutoTokenizer.from_pretrained(self.tokenizer_name, legacy=False, **kwargs)
         if self.tokenizer.chat_template is None:
             self.eos += extra_eos_for_direct_completion(dataset)
 
@@ -242,7 +248,7 @@ def __init__(self, name: str, **kwargs):
         super().__init__(name=name, **kwargs)
         self.eos += ["\n```\n"]
         print(f"EOS strings: {self.eos}")
-        self.tokenizer = AutoTokenizer.from_pretrained(self.name, legacy=False, **kwargs)
+        self.tokenizer = AutoTokenizer.from_pretrained(self.tokenizer_name if self.tokenizer_name else self.name, **kwargs)
 
     def codegen(
         self, prompt: str, do_sample: bool = True, num_samples: int = 200
@@ -475,6 +481,7 @@ def make_model(
     tp=1,
     base_url=None,
     trust_remote_code=False,
+    tokenizer_name=None,
 ):
     if backend == "vllm":
         return GeneralVllmDecoder(
@@ -484,6 +491,7 @@ def make_model(
             dataset=dataset,
             tp=tp,
             trust_remote_code=trust_remote_code,
+            tokenizer_name=tokenizer_name,
         )
     elif backend == "hf":
         return GenenralHfTorchDecoder(
@@ -492,6 +500,7 @@ def make_model(
             temperature=temperature,
             dataset=dataset,
             trust_remote_code=trust_remote_code,
+            tokenizer_name=tokenizer_name,
         )
     elif backend == "openai":
         return OpenAIChatDecoder(