Skip to content

Commit

Permalink
Multithread api call in Wildbench
Browse files Browse the repository at this point in the history
  • Loading branch information
jmercat committed Nov 20, 2024
1 parent 757bf3b commit b7b27bb
Showing 1 changed file with 15 additions and 4 deletions.
19 changes: 15 additions & 4 deletions eval/chat_benchmarks/WildBench/eval_instruct.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
import jsonlines
from datasets import load_dataset
from openai import OpenAI
from concurrent.futures import ThreadPoolExecutor, as_completed

from lm_eval.api.instance import Instance
from lm_eval.api.model import LM
Expand Down Expand Up @@ -49,6 +50,7 @@ class WildBenchConfig:
model: str = "gpt-4o-mini-2024-07-18"
mode: str = "score"
batch_mode: bool = True
api_parallel: int = 32

# Task weights
task_weights: Dict[str, float] = None
Expand Down Expand Up @@ -336,15 +338,24 @@ def _process_evaluator_file(self, eval_file: str, client: OpenAI) -> None:
with open(eval_file, "r") as file:
lines = file.readlines()

results = []
for line in lines:
def process_line(line):
payload = json.loads(line)
payload["body"]["max_tokens"] = 4096
response = client.chat.completions.create(**payload["body"])

result = payload.copy()
result["response"] = json.loads(response.json())
results.append(result)
return result

results = []
# Use ThreadPoolExecutor since API calls are I/O bound
with ThreadPoolExecutor(max_workers=self.config.api_parallel) as executor:
future_to_line = {executor.submit(process_line, line): line for line in lines}
for future in as_completed(future_to_line):
try:
result = future.result()
results.append(result)
except Exception as e:
self.logger.error(f"Error processing line: {str(e)}")

output_file = eval_file.replace("batch-submit.jsonl", "results.jsonl")
with open(output_file, "w") as file:
Expand Down

0 comments on commit b7b27bb

Please sign in to comment.