add custom tokenzier

Author: marianna13

bigcodebench/model.py

@@ -92,6 +92,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
@@ -102,6 +103,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(
@@ -129,11 +131,13 @@ 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, **kwargs)
+        if self.tokenizer_name is None:
+            self.tokenizer_name = self.name
+        
+        self.tokenizer = AutoTokenizer.from_pretrained(self.tokenizer_name, **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)
+        self.llm = LLM(model=name, max_model_len=2048, tokenizer=self.tokenizer_name, **kwargs)
 
     def is_direct_completion(self) -> bool:
         return self.tokenizer.chat_template is None
@@ -185,9 +189,12 @@ def __init__(self, name: str, dataset: str, **kwargs):
         kwargs["torch_dtype"] = getattr(torch, self.dtype)
         self.skip_special_tokens = True
 
-        print(f"{kwargs = }")
+        print(f"{kwargs = }", self.tokenizer_name)
+
+        if self.tokenizer_name is None:
+            self.tokenizer_name = self.name
 
-        self.tokenizer = AutoTokenizer.from_pretrained(name, **kwargs)
+        self.tokenizer = AutoTokenizer.from_pretrained(self.tokenizer_name, **kwargs)
         if self.tokenizer.chat_template is None:
             self.eos += extra_eos_for_direct_completion(dataset)
 
@@ -253,7 +260,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, **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
@@ -486,6 +493,7 @@ def make_model(
     tp=1,
     base_url=None,
     trust_remote_code=False,
+    tokenizer_name=None,
 ):
     if backend == "vllm":
         return GeneralVllmDecoder(
@@ -495,6 +503,7 @@ def make_model(
             dataset=dataset,
             tp=tp,
             trust_remote_code=trust_remote_code,
+            tokenizer_name=tokenizer_name,
         )
     elif backend == "hf":
         return GenenralHfTorchDecoder(
@@ -503,6 +512,7 @@ def make_model(
             temperature=temperature,
             dataset=dataset,
             trust_remote_code=trust_remote_code,
+            tokenizer_name=tokenizer_name,
         )
     elif backend == "openai":
         return OpenAIChatDecoder(