Optimize evaluation performance using AsyncOpenAI for concurrent processing

Author: sahilds1

evaluation/evals.py

@@ -23,6 +23,7 @@
 
 import argparse
 import logging
+import asyncio
 
 import pandas as pd
 
@@ -37,15 +38,15 @@
 )
 
 
-def evaluate_response(model: str, instructions: str, input: str) -> pd.DataFrame:
+async def evaluate_response(model: str, instructions: str, input: str) -> pd.DataFrame:
     """
     Test a prompt with a set of test data by scoring each item in the data set
     """
 
     try:
         handler = ModelFactory.get_handler(model)
 
-        generated_text, token_usage, pricing, duration = handler.handle_request(
+        generated_text, token_usage, pricing, duration = await handler.handle_request(
             instructions, input
         )
 
@@ -145,7 +146,7 @@ def load_csv(file_path: str, required_columns: list) -> pd.DataFrame:
     return df
 
 
-if __name__ == "__main__":
+async def main():
     # TODO: Add test evaluation argument to run on the first 10 rows of the dataset file
 
     parser = argparse.ArgumentParser()
@@ -177,14 +178,15 @@ def load_csv(file_path: str, required_columns: list) -> pd.DataFrame:
     # Bulk model and prompt experimentation: Cross join the experiment and dataset DataFrames
     df_in = df_experiment.merge(df_dataset, how="cross")
 
-    # Evaluate each row in the input DataFrame
-    results = []
-    for index, row in enumerate(df_in.itertuples(index=False)):
-        result = evaluate_response(row.MODEL, row.INSTRUCTIONS, row.INPUT)
-        results.append(result)
+    # Evaluate each row in the input DataFrame concurrently
+    logging.info(f"Starting evaluation of {len(df_in)} rows")
+    tasks = [
+        evaluate_response(row.MODEL, row.INSTRUCTIONS, row.INPUT)
+        for row in df_in.itertuples(index=False)
+    ]
 
-        # TODO: Use tqdm or similar library to show progress bar
-        logging.info(f"Processed row {index + 1}/{len(df_in)}")
+    results = await asyncio.gather(*tasks)
+    logging.info(f"Completed evaluation of {len(results)} rows")
 
     df_evals = pd.concat(results, axis=0, ignore_index=True)
 
@@ -195,3 +197,7 @@ def load_csv(file_path: str, required_columns: list) -> pd.DataFrame:
     df_out.to_csv(args.results, index=False)
     logging.info(f"Results saved to {args.results}")
     logging.info("Evaluation completed successfully.")
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

server/api/services/llm_services.py

@@ -7,12 +7,12 @@
 import logging
 from abc import ABC, abstractmethod
 
-import openai
+from openai import AsyncOpenAI
 
 
 class BaseModelHandler(ABC):
     @abstractmethod
-    def handle_request(
+    async def handle_request(
         self, query: str, context: str
     ) -> tuple[str, dict[str, int], dict[str, float], float]:
         pass
@@ -31,9 +31,9 @@ class GPT4OMiniHandler(BaseModelHandler):
     PRICING_DOLLARS_PER_MILLION_TOKENS = {"input": 0.15, "output": 0.60}
 
     def __init__(self) -> None:
-        self.client = openai.OpenAI(api_key=os.environ.get("OPENAI_API_KEY"))
+        self.client = AsyncOpenAI(api_key=os.environ.get("OPENAI_API_KEY"))
 
-    def handle_request(
+    async def handle_request(
         self, query: str, context: str
     ) -> tuple[str, dict[str, int], dict[str, float], float]:
         """
@@ -46,7 +46,7 @@ def handle_request(
         """
         start_time = time.time()
         # TODO: Add error handling for API requests and invalid responses
-        response = self.client.responses.create(
+        response = await self.client.responses.create(
             model=self.MODEL, instructions=query, input=context, temperature=0.0
         )
         duration = time.time() - start_time
@@ -123,9 +123,9 @@ class GPT41NanoHandler(BaseModelHandler):
     """
 
     def __init__(self) -> None:
-        self.client = openai.OpenAI(api_key=os.environ.get("OPENAI_API_KEY"))
+        self.client = AsyncOpenAI(api_key=os.environ.get("OPENAI_API_KEY"))
 
-    def handle_request(
+    async def handle_request(
         self, query: str, context: str
     ) -> tuple[str, dict[str, int], dict[str, float], float]:
         """
@@ -144,7 +144,7 @@ def handle_request(
         start_time = time.time()
         # TODO: Add error handling for API requests and invalid responses
 
-        response = self.client.responses.create(
+        response = await self.client.responses.create(
             model=self.MODEL, instructions=query, input=context, temperature=0.0
         )
         duration = time.time() - start_time