feat: batch o1 and deepseek-chat via concurrency

Author: terryyz

bigcodebench/gen/util/openai_request.py

@@ -1,4 +1,3 @@
-import signal
 import time
 
 import openai
@@ -14,53 +13,38 @@ def make_request(
     n: int = 1,
     **kwargs
 ) -> ChatCompletion:
-    system_msg = "You are a helpful assistant good at coding."
-    if (
-        kwargs.get("response_format", None)
-        and kwargs["response_format"]["type"] == "json_object"
-    ):
-        system_msg = "You are a helpful assistant designed to output JSON."
-
+    kwargs["top_p"] = 0.95
+    kwargs["max_completion_tokens"] = max_tokens
+    if model.startswith("o1-"):  # pop top-p and max_completion_tokens
+        kwargs.pop("top_p")
+        kwargs.pop("max_completion_tokens")
+    
     return client.chat.completions.create(
         model=model,
         messages=[
-            {"role": "system", "content": system_msg},
             {"role": "user", "content": message},
         ],
-        max_tokens=max_tokens,
         temperature=temperature,
         n=n,
         **kwargs
     )
 
 
-def handler(signum, frame):
-    # swallow signum and frame
-    raise Exception("end of time")
-
-
 def make_auto_request(*args, **kwargs) -> ChatCompletion:
     ret = None
     while ret is None:
         try:
-            signal.signal(signal.SIGALRM, handler)
-            signal.alarm(100)
             ret = make_request(*args, **kwargs)
-            signal.alarm(0)
         except openai.RateLimitError:
             print("Rate limit exceeded. Waiting...")
-            signal.alarm(0)
             time.sleep(5)
         except openai.APIConnectionError:
             print("API connection error. Waiting...")
-            signal.alarm(0)
             time.sleep(5)
         except openai.APIError as e:
             print(e)
-            signal.alarm(0)
         except Exception as e:
             print("Unknown error. Waiting...")
             print(e)
-            signal.alarm(0)
             time.sleep(1)
-    return ret
+    return ret

bigcodebench/provider/openai.py

@@ -1,12 +1,12 @@
 import os
 from typing import List
-from tqdm import tqdm
 
 import openai
 
-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
+from bigcodebench.provider.base import DecoderBase
+from bigcodebench.provider.utility import concurrent_call
 
 class OpenAIChatDecoder(DecoderBase):
     def __init__(self, name: str, base_url=None, **kwargs) -> None:
@@ -15,34 +15,83 @@ def __init__(self, name: str, base_url=None, **kwargs) -> None:
             api_key=os.getenv("OPENAI_API_KEY", "none"), base_url=base_url
         )
 
+    # def codegen(
+    #     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"
+    #     all_outputs = []
+    #     for prompt in tqdm(prompts):
+    #         outputs = []
+    #         message = make_raw_chat_prompt(
+    #             task_prompt=prompt,
+    #             subset=self.subset,
+    #             split=self.split,
+    #             instruction_prefix=self.instruction_prefix,
+    #             response_prefix=self.response_prefix,
+    #             tokenizer=None,
+    #         )
+    #         ret = make_auto_request(
+    #             self.client,
+    #             message=message,
+    #             model=self.name,
+    #             max_tokens=self.max_new_tokens,
+    #             temperature=self.temperature,
+    #             n=num_samples,
+    #         )
+    #         for item in ret.choices:
+    #             outputs.append(item.message.content)
+    #         all_outputs.append(outputs)
+    #     return all_outputs
+
+    # def is_direct_completion(self) -> bool:
+    #     return False
+    
     def codegen(
         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"
+        messages = [make_raw_chat_prompt(
+            task_prompt=prompt,
+            subset=self.subset,
+            split=self.split,
+            instruction_prefix=self.instruction_prefix,
+            response_prefix=self.response_prefix,
+            tokenizer=None,
+        ) for prompt in prompts]
+        # use concurrency based batching for o1 and deepseek models
+        if self.name.startswith("o1-") or self.name == "deepseek-chat":
+            return self._codegen_batch_via_concurrency(messages, num_samples)
+
+        return self._codegen_api_batch(messages, num_samples)
+
+    def _codegen_api_batch(self, messages: List[str], num_samples: int) -> List[str]:
+        client = openai.OpenAI(
+            api_key=os.getenv("OPENAI_API_KEY", "none"), base_url=self.base_url
+        )
+        
         all_outputs = []
-        for prompt in tqdm(prompts):
-            outputs = []
-            message = make_raw_chat_prompt(
-                task_prompt=prompt,
-                subset=self.subset,
-                split=self.split,
-                instruction_prefix=self.instruction_prefix,
-                response_prefix=self.response_prefix,
-                tokenizer=None,
-            )
+        for message in messages:
             ret = make_auto_request(
-                self.client,
+                client,
                 message=message,
                 model=self.name,
                 max_tokens=self.max_new_tokens,
                 temperature=self.temperature,
                 n=num_samples,
             )
+            outputs = []
             for item in ret.choices:
                 outputs.append(item.message.content)
             all_outputs.append(outputs)
         return all_outputs
 
+    def _codegen_batch_via_concurrency(self, messages: List[str], num_samples: int) -> List[str]:
+        batches = concurrent_call(
+            num_samples, self._codegen_api_batch, messages, num_samples=1
+        )
+        return [b[0] for b in batches]
+
     def is_direct_completion(self) -> bool:
         return False

bigcodebench/provider/utility.py

@@ -1,5 +1,6 @@
 from typing import List
 from transformers import AutoTokenizer
+from concurrent.futures import ThreadPoolExecutor
 
 EOS = [
     "<|endoftext|>",
@@ -64,4 +65,10 @@ def make_raw_chat_prompt(
             ],
             tokenize=False,
         ).split(_MAGIC_SPLITTER_)[0]
-    return task_prompt
+    return task_prompt
+
+
+def concurrent_call(n, callback, /, *args, **kwargs):
+    with ThreadPoolExecutor(max_workers=n) as executor:
+        futures = [executor.submit(callback, *args, **kwargs) for _ in range(n)]
+        return [future.result() for future in futures]