feat: make prefill optional

Author: terryyz

bigcodebench/generate.py

@@ -140,8 +140,9 @@ def run_codegen(
     backend: str = "vllm",
     base_url: str = None,
     tp: int = 1,
-    instruction_prefix: str = None,
-    response_prefix: str = None,
+    instruction_prefix: str = "Please provide a self-contained Python script that solves the following problem in a markdown code block:",
+    response_prefix: str ="Below is a Python script with a self-contained function that solves the problem and passes corresponding tests:",
+    prefill: bool = True,
     revision: str = "main",
     trust_remote_code: bool = False,
     tokenizer_name: str = None,
@@ -163,11 +164,6 @@ def run_codegen(
     # Make project dir
     os.makedirs(root, exist_ok=True)
 
-    if instruction_prefix is None:  
-        instruction_prefix = "Please provide a self-contained Python script that solves the following problem in a markdown code block:"
-    if response_prefix is None:
-        response_prefix = "Below is a Python script with a self-contained function that solves the problem and passes corresponding tests:"
-    
     # Make dir for codes generated by each model
     model_runner = make_model(
         model=model,
@@ -179,6 +175,7 @@ def run_codegen(
         reasoning_effort=reasoning_effort,
         instruction_prefix=instruction_prefix,
         response_prefix=response_prefix,
+        prefill=prefill,
         base_url=base_url,
         tp=tp,
         revision=revision,

bigcodebench/provider/__init__.py

@@ -14,6 +14,7 @@ def make_model(
     # instruction model only
     instruction_prefix: str = None,
     response_prefix: str = None,
+    prefill: bool = True,
     # vllm and hf only
     revision: str = "main",
     # vllm only
@@ -42,6 +43,7 @@ def make_model(
             tp=tp,
             instruction_prefix=instruction_prefix,
             response_prefix=response_prefix,
+            prefill=prefill,
             trust_remote_code=trust_remote_code,
             tokenizer_name=tokenizer_name,
             tokenizer_legacy=tokenizer_legacy,
@@ -60,6 +62,7 @@ def make_model(
             direct_completion=direct_completion,
             instruction_prefix=instruction_prefix,
             response_prefix=response_prefix,
+            prefill=prefill,
             attn_implementation=attn_implementation,
             trust_remote_code=trust_remote_code,
             tokenizer_name=tokenizer_name,

bigcodebench/provider/base.py

@@ -20,6 +20,7 @@ def __init__(
         tokenizer_legacy: bool = False,
         instruction_prefix: str = None,
         response_prefix: str = None,
+        prefill: bool = True,
     ) -> None:
         print("Initializing a decoder model: {} ...".format(name))
         self.name = name
@@ -37,6 +38,7 @@ def __init__(
         self.tokenizer_legacy = tokenizer_legacy
         self.instruction_prefix = instruction_prefix
         self.response_prefix = response_prefix
+        self.prefill = prefill
 
     @abstractmethod
     def codegen(

bigcodebench/provider/utility.py

@@ -28,6 +28,7 @@ def make_raw_chat_prompt(
     split: str, 
     instruction_prefix: str,
     response_prefix: str,
+    prefill: bool,
     tokenizer: AutoTokenizer,
     direct_completion: bool = False,
 ) -> str:
@@ -58,13 +59,21 @@ def make_raw_chat_prompt(
 ```
 """
     if tokenizer:
-        task_prompt = tokenizer.apply_chat_template(
-            [
-                {"role": "user", "content": task_prompt},
-                {"role": "assistant", "content": response},
-            ],
-            tokenize=False,
-        ).split(_MAGIC_SPLITTER_)[0]
+        if prefill:
+            task_prompt = tokenizer.apply_chat_template(
+                [
+                    {"role": "user", "content": task_prompt},
+                    {"role": "assistant", "content": response},
+                ],
+                tokenize=False,
+            ).split(_MAGIC_SPLITTER_)[0]
+        else:
+            task_prompt = tokenizer.apply_chat_template(
+                [
+                    {"role": "user", "content": task_prompt},
+                ],
+                tokenize=False,
+            ).split(_MAGIC_SPLITTER_)[0]
     return task_prompt
 
 

bigcodebench/provider/vllm.py

@@ -47,6 +47,7 @@ def codegen(
                 split=self.split,
                 instruction_prefix=self.instruction_prefix,
                 response_prefix=self.response_prefix,
+                prefill=self.prefill,
                 tokenizer=self.tokenizer,
                 direct_completion=self.direct_completion,
             )