run single turn rollouts in parallel

Author: Dylan Huang

eval_protocol/pytest/default_single_turn_rollout_process.py

@@ -1,36 +1,42 @@
+import asyncio
 from typing import List
 
-from openai import OpenAI
+from openai import AsyncOpenAI
 
 from eval_protocol.auth import get_fireworks_api_base, get_fireworks_api_key
-from eval_protocol.models import CompletionParams, EvaluationRow, InputMetadata, Message
-from eval_protocol.pytest.types import Dataset, ModelParam, RolloutProcessorConfig
+from eval_protocol.models import EvaluationRow, Message
+from eval_protocol.pytest.types import RolloutProcessorConfig
 
 
-def default_single_turn_rollout_processor(
+async def default_single_turn_rollout_processor(
     rows: List[EvaluationRow], config: RolloutProcessorConfig
 ) -> List[EvaluationRow]:
-    """Generate a single response from a Fireworks model."""
+    """Generate a single response from a Fireworks model concurrently."""
 
     api_key = get_fireworks_api_key()
     api_base = get_fireworks_api_base()
-    client = OpenAI(api_key=api_key, base_url=f"{api_base}/inference/v1")
+    client = AsyncOpenAI(api_key=api_key, base_url=f"{api_base}/inference/v1")
 
-    dataset: Dataset = []
-    for row in rows:
+    async def process_row(row: EvaluationRow) -> EvaluationRow:
+        """Process a single row asynchronously."""
         if len(row.messages) == 0:
             raise ValueError("Messages is empty. Please provide a non-empty dataset")
 
         messages_payload = [{"role": m.role, "content": m.content} for m in row.messages]
 
-        response = client.chat.completions.create(model=config.model, messages=messages_payload, **config.input_params)
+        response = await client.chat.completions.create(
+            model=config.model, messages=messages_payload, **config.input_params
+        )
         assistant_content = response.choices[0].message.content or ""
         messages = list(row.messages) + [Message(role="assistant", content=assistant_content)]
-        processed = EvaluationRow(
+
+        return EvaluationRow(
             messages=messages,
-            ground_truth=row.ground_truth,
-            input_metadata=InputMetadata(completion_params=CompletionParams(model=config.model)),
+            **row.model_dump(exclude={"messages"}),
         )
 
-        dataset.append(processed)
+    # Process all rows concurrently
+    tasks = [process_row(row) for row in rows]
+    dataset = await asyncio.gather(*tasks)
+
     return dataset