refactor(gen): update model provider

Author: terryyz

bigcodebench/provider/__init__.py

@@ -8,19 +8,20 @@ def make_model(
     split: str,
     dataset: str = "bigcodebench",
     temperature: float = 0.0,
+    max_new_tokens: int = 1280,
     # instruction model only
-    instruction_prefix=None,
-    response_prefix=None,
+    instruction_prefix: str = None,
+    response_prefix: str = None,
     # vllm only
-    tp=1,
-    direct_completion=False,
-    base_url=None,
-    trust_remote_code=False,
+    tp: int = 1,
+    direct_completion: bool = False,
+    base_url: str = None,
+    trust_remote_code: bool = False,
     # hf only
-    attn_implementation="eager",
+    attn_implementation: str = "eager",
     # tokenizer
-    tokenizer_name=None,
-    tokenizer_kwargs=None,
+    tokenizer_name: str = None,
+    tokenizer_legacy: bool = True,
 ) -> DecoderBase:
     if backend == "vllm":
         from bigcodebench.provider.vllm import VllmDecoder
@@ -30,9 +31,10 @@ def make_model(
             subset=subset,
             split=split,
             temperature=temperature,
+            max_new_tokens=max_new_tokens,
             dataset=dataset,
             direct_completion=direct_completion,
-            tensor_parallel_size=tp,
+            tp=tp,
             instruction_prefix=instruction_prefix,
             response_prefix=response_prefix,
         )
@@ -44,6 +46,7 @@ def make_model(
             subset=subset,
             split=split,
             temperature=temperature,
+            max_new_tokens=max_new_tokens,
             dataset=dataset,
             direct_completion=direct_completion,
             instruction_prefix=instruction_prefix,
@@ -59,6 +62,7 @@ def make_model(
             subset=subset,
             split=split,
             temperature=temperature,
+            max_new_tokens=max_new_tokens,
             base_url=base_url,
             instruction_prefix=instruction_prefix,
             response_prefix=response_prefix,
@@ -71,6 +75,7 @@ def make_model(
             subset=subset,
             split=split,
             temperature=temperature,
+            max_new_tokens=max_new_tokens,
             instruction_prefix=instruction_prefix,
             response_prefix=response_prefix,
         )
@@ -83,6 +88,7 @@ def make_model(
             subset=subset,
             split=split,
             temperature=temperature,
+            max_new_tokens=max_new_tokens,
             instruction_prefix=instruction_prefix,
             response_prefix=response_prefix,
         )
@@ -95,6 +101,7 @@ def make_model(
             subset=subset,
             split=split,
             temperature=temperature,
+            max_new_tokens=max_new_tokens,
             instruction_prefix=instruction_prefix,
             response_prefix=response_prefix,
         )

bigcodebench/provider/anthropic.py

@@ -1,9 +1,10 @@
 import os
 from typing import List
+from tqdm import tqdm
 
 import anthropic
 
-from bigcodebench.gen.util import anthropic_request
+from bigcodebench.gen.util.anthropic_request import make_auto_request
 from bigcodebench.provider.base import DecoderBase
 from bigcodebench.provider.utility import make_raw_chat_prompt
 
@@ -18,15 +19,12 @@ def codegen(
         if do_sample:
             assert self.temperature > 0, "Temperature must be positive for sampling"
 
-        if not do_sample:
-            assert batch_size == 1, "Sampling only supports batch size of 1"
-
         all_outputs = []
         for prompt in tqdm(prompts):
             outputs = []
 
             for _ in range(num_samples):
-                message = anthropic_request.make_auto_request(
+                ret = make_auto_request(
                     client=self.client,
                     model=self.name,
                     messages=[
@@ -46,9 +44,9 @@ def codegen(
                     temperature=self.temperature,
                     stop_sequences=self.eos,
                 )
-                outputs.append(message.content[0].text)
+                outputs.append(ret.content[0].text)
             all_outputs.append(outputs)
-        return outputs
+        return all_outputs
 
     def is_direct_completion(self) -> bool:
         return False

bigcodebench/provider/base.py

@@ -1,7 +1,7 @@
 from abc import ABC, abstractmethod
 from typing import List
 
-from evalplus.provider.utility import EOS
+from bigcodebench.provider.utility import EOS
 
 
 class DecoderBase(ABC):
@@ -11,12 +11,14 @@ def __init__(
         subset: str,
         split: str,
         temperature: float = 0.8,
-        max_new_tokens: int = 5120,
+        max_new_tokens: int = 1280,
         dtype: str = "bfloat16",  # default
         direct_completion: bool = False,
         trust_remote_code: bool = False,
         tokenizer_name: str = None,
         tokenizer_legacy: bool = False,
+        instruction_prefix: str = None,
+        response_prefix: str = None,
     ) -> None:
         print("Initializing a decoder model: {} ...".format(name))
         self.name = name
@@ -31,6 +33,8 @@ def __init__(
         self.trust_remote_code = trust_remote_code
         self.tokenizer_name = tokenizer_name
         self.tokenizer_legacy = tokenizer_legacy
+        self.instruction_prefix = instruction_prefix
+        self.response_prefix = response_prefix
 
     @abstractmethod
     def codegen(

bigcodebench/provider/google.py

@@ -1,9 +1,9 @@
 import os
 from typing import List
+from tqdm import tqdm
 
 import google.generativeai as genai
 
-
 from bigcodebench.provider.base import DecoderBase
 from bigcodebench.gen.util.google_request import make_auto_request
 from bigcodebench.provider.utility import make_raw_chat_prompt
@@ -24,7 +24,7 @@ def codegen(
         all_outputs = []
 
         for prompt in tqdm(prompts):
-            ret_texts = []
+            outputs = []
             message = make_raw_chat_prompt(
                 task_prompt=prompt,
                 subset=self.subset,
@@ -33,25 +33,23 @@ def codegen(
                 response_prefix=self.response_prefix,
                 tokenizer=None,
             )
-            replies = make_auto_request(
+            ret = make_auto_request(
                 self.client,
                 message,
                 self.name,
-                n=batch_size,
+                n=num_samples,
                 max_tokens=self.max_new_tokens,
                 temperature=self.temperature,
             )
-            for candidate in replies.candidates:
+            for candidate in ret.candidates:
                 parts = candidate.content.parts
                 if parts:
-                    ret_texts.append(parts[0].text)
+                    outputs.append(parts[0].text)
                 else:
                     print("Empty response!")
-                    ret_texts.append("")
+                    outputs.append("")
                     print(f"{candidate.safety_ratings = }")
-            ret_texts.append("")
-            all_outputs.append(ret_texts + [""] * (batch_size - len(ret_texts)))
-
+            all_outputs.append(outputs)
         return all_outputs
 
     def is_direct_completion(self) -> bool:

bigcodebench/provider/hf.py

@@ -4,8 +4,8 @@
 from stop_sequencer import StopSequencer
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
-from evalplus.provider.base import DecoderBase
-from evalplus.provider.utility import (
+from bigcodebench.provider.base import DecoderBase
+from bigcodebench.provider.utility import (
     extra_eos_for_direct_completion,
     make_raw_chat_prompt,
 )
@@ -33,14 +33,17 @@ def __init__(
         print(f"{kwargs = }")
 
         self.tokenizer = AutoTokenizer.from_pretrained(name, use_fast=False, legacy=self.tokenizer_legacy)
+        self.tokenizer.pad_token = self.tokenizer.eos_token
+        # assume the model is decoder-only
+        self.tokenizer.padding_side = 'left'
+        
         if self.is_direct_completion():  # no chat template
             self.eos += extra_eos_for_direct_completion(dataset)
         else:  # with chat template
             self.eos += ["\n```\n"]
 
         print(f"{self.eos = }")
         self.model = AutoModelForCausalLM.from_pretrained(name, **kwargs)
-        self.model = self.model.to(self.device)
 
     def is_direct_completion(self) -> bool:
         return self.direct_completion or self.tokenizer.chat_template is None
@@ -61,15 +64,16 @@ def codegen(
             )
             for prompt in prompts
         ]
-        input_tokens = self.tokenizer.encode(prompts, return_tensors="pt").to(
+        
+        input_tokens = self.tokenizer(prompts, return_tensors="pt", padding=True, truncation=True).to(
             self.device
-        )
+        )["input_ids"]
+        
         kwargs = {}
         if do_sample:
             kwargs["top_p"] = 0.95
             kwargs["temperature"] = self.temperature
-
-        outputs = self.model.generate(
+        ret = self.model.generate(
             input_tokens,
             max_new_tokens=self.max_new_tokens,
             do_sample=do_sample,
@@ -79,17 +83,23 @@ def codegen(
             tokenizer=self.tokenizer,
             **kwargs,
         )
+        
+        # Reshape ret into a list of lists, each sublist containing num_samples elements
+        ret_chunks = [ret[i:i + num_samples] for i in range(0, len(ret), num_samples)]
 
-        gen_strs = self.tokenizer.batch_decode(
-            outputs[:, input_tokens.size(-1) :],
-            skip_special_tokens=self.skip_special_tokens,
-        )
-        outputs = []
-        # removes eos tokens.
-        for output in gen_strs:
-            min_index = 10000
-            for eos in self.eos:
-                if eos in output:
-                    min_index = min(min_index, output.index(eos))
-            outputs.append(output[:min_index].replace("\t", "    "))
-        return outputs
+        all_outputs = []
+        # Process each chunk in ret_chunks
+        for i, ret_chunk in enumerate(ret_chunks):
+            gen_strs = self.tokenizer.batch_decode(
+                ret_chunk[:, input_tokens[i].size(-1):],
+                skip_special_tokens=self.skip_special_tokens,
+            )
+            outputs = []
+            for output in gen_strs:
+                min_index = 10000
+                for eos in self.eos:
+                    if eos in output:
+                        min_index = min(min_index, output.index(eos))
+                outputs.append(output[:min_index].replace("\t", "    "))
+            all_outputs.append(outputs)
+        return all_outputs

bigcodebench/provider/mistral.py

@@ -1,18 +1,21 @@
 import os
 from typing import List
+from tqdm import tqdm
 
-import anthropic
+from mistralai.client import MistralClient
+from mistralai.models.chat_completion import ChatMessage
 
 from bigcodebench.provider.base import DecoderBase
+from bigcodebench.gen.util.mistral_request import make_auto_request
 from bigcodebench.provider.utility import make_raw_chat_prompt
 
-class MistralDecoder(DecoderBase):
+class MistralChatDecoder(DecoderBase):
     def __init__(self, name: str, **kwargs) -> None:
         super().__init__(name, **kwargs)
-        self.client = mistral.Mistral(api_key=os.getenv("MISTRAL_API_KEY"))
+        self.client = MistralClient(api_key=os.getenv("MISTRAL_API_KEY"))
 
     def codegen(
-        self, prompt: str, do_sample: bool = True, num_samples: int = 200
+        self, prompts: List[str], do_sample: bool = True, num_samples: int = 200
     ) -> List[str]:
         if do_sample:
             assert self.temperature > 0, "Temperature must be positive for sampling"
@@ -22,7 +25,7 @@ def codegen(
             outputs = []
 
             for _ in range(num_samples):
-                message = mistral_request.make_auto_request(
+                ret = make_auto_request(
                     client=self.client,
                     model=self.name,
                     messages=[
@@ -40,9 +43,8 @@ def codegen(
                         )
                     ],
                     max_tokens=self.max_new_tokens,
-                    **kwargs,
                 )
-                outputs.append(message.content[0].text)
+                outputs.append(ret.choices[0].message.content)
             all_outputs.append(outputs)
         return all_outputs
 

bigcodebench/provider/openai.py

@@ -1,10 +1,11 @@
 import os
 from typing import List
+from tqdm import tqdm
 
 import openai
 
-from evalplus.gen.util import openai_request
-from evalplus.provider.base import DecoderBase
+from bigcodebench.provider.base import DecoderBase
+from bigcodebench.gen.util.openai_request import make_auto_request
 from bigcodebench.provider.utility import make_raw_chat_prompt
 
 class OpenAIChatDecoder(DecoderBase):
@@ -21,6 +22,7 @@ def codegen(
             assert self.temperature > 0, "Temperature must be positive for sampling"
         all_outputs = []
         for prompt in tqdm(prompts):
+            outputs = []
             message = make_raw_chat_prompt(
                 task_prompt=prompt,
                 subset=self.subset,
@@ -29,7 +31,7 @@ def codegen(
                 response_prefix=self.response_prefix,
                 tokenizer=None,
             )
-            ret = openai_request.make_auto_request(
+            ret = make_auto_request(
                 self.client,
                 message=message,
                 model=self.name,