Refactor evals.py for clarity, formatting, and maintainability

Author: sahilds1

server/api/services/evals.py

@@ -2,7 +2,7 @@
 Evaluate LLM outputs using multiple metrics and compute associated costs
 """
 
-#TODO: Add tests on a small dummy dataset to confirm it handles errors gracefully and produces expected outputs
+# TODO: Add tests on a small dummy dataset to confirm it handles errors gracefully and produces expected outputs
 
 import argparse
 import logging
@@ -13,11 +13,14 @@
 
 from services import ModelFactory
 
-logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+logging.basicConfig(
+    level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
+)
 
 
-#TODO: Rename model to model name and query to instructions for clarity
-def evaluate_response(model_name: str, query: str, context: str, reference: str) -> pd.DataFrame:
+def evaluate_response(
+    model_name: str, query: str, context: str, reference: str
+) -> pd.DataFrame:
     """
     Evaluates the response of a model to a given query and context, computes extractiveness metrics, token usage, and cost
 
@@ -33,35 +36,47 @@ def evaluate_response(model_name: str, query: str, context: str, reference: str)
 
     handler = ModelFactory.get_handler(model_name)
 
-    #TODO: Add error handling for unsupported models
-        
+    # TODO: Add error handling for unsupported models
+
     output_text, token_usage, pricing, duration = handler.handle_request(query, context)
 
     doc = Doc(query="", choices=[], gold_index=0, specific={"text": context})
-    extractiveness = Extractiveness().compute(formatted_doc=doc, predictions=[output_text])
+    extractiveness = Extractiveness().compute(
+        formatted_doc=doc, predictions=[output_text]
+    )
 
-    input_cost_dollars = (pricing['input'] / 1000000) * token_usage.input_tokens
-    output_cost_dollars = (pricing['output'] / 1000000) * token_usage.output_tokens
+    input_cost_dollars = (pricing["input"] / 1000000) * token_usage.input_tokens
+    output_cost_dollars = (pricing["output"] / 1000000) * token_usage.output_tokens
 
     total_cost_dollars = input_cost_dollars + output_cost_dollars
 
-    return pd.DataFrame([{
-            "Output Text": output_text,
-            "Extractiveness Coverage": extractiveness['summarization_coverage'],
-            "Extractiveness Density": extractiveness['summarization_density'],
-            "Extractiveness Compression": extractiveness['summarization_compression'],
-            "Input Token Usage": token_usage.input_tokens,
-            "Output Token Usage": token_usage.output_tokens,
-            "Cost (USD)": total_cost_dollars,
-            "Duration (s)": duration
-        }])
+    return pd.DataFrame(
+        [
+            {
+                "Output Text": output_text,
+                "Extractiveness Coverage": extractiveness["summarization_coverage"],
+                "Extractiveness Density": extractiveness["summarization_density"],
+                "Extractiveness Compression": extractiveness[
+                    "summarization_compression"
+                ],
+                "Input Token Usage": token_usage.input_tokens,
+                "Output Token Usage": token_usage.output_tokens,
+                "Cost (USD)": total_cost_dollars,
+                "Duration (s)": duration,
+            }
+        ]
+    )
 
 
 if __name__ == "__main__":
-
-    parser = argparse.ArgumentParser(description="Evaluate LLM outputs using multiple metrics and compute associated costs")
+    # TODO: Add CLI argument to specify the metrics to be computed
+    parser = argparse.ArgumentParser(
+        description="Evaluate LLM outputs using multiple metrics and compute associated costs"
+    )
     parser.add_argument("--config", "-c", required=True, help="Path to config CSV file")
-    parser.add_argument("--reference", "-r", required=True, help="Path to reference CSV file")
+    parser.add_argument(
+        "--reference", "-r", required=True, help="Path to reference CSV file"
+    )
     parser.add_argument("--output", "-o", required=True, help="Path to output CSV file")
 
     args = parser.parse_args()
@@ -73,34 +88,46 @@ def evaluate_response(model_name: str, query: str, context: str, reference: str)
     # Remove the trailing whitespace from column names
     df_config.columns = df_config.columns.str.strip()
 
-    #TODO: Check if the required columns are present
+    # TODO: Check if the required columns are present
 
     # Check if all models in the config are supported by ModelFactory
-    if not all(model in ModelFactory.HANDLERS.keys() for model in df_config['Model'].unique()):
-        raise ValueError(f"Unsupported model(s) found in config: {set(df_config['Model'].unique()) - set(ModelFactory.HANDLERS.keys())}")
-    
+    if not all(
+        model in ModelFactory.HANDLERS.keys()
+        for model in df_config["Model Name"].unique()
+    ):
+        raise ValueError(
+            f"Unsupported model(s) found in config: {set(df_config['Model Name'].unique()) - set(ModelFactory.HANDLERS.keys())}"
+        )
+
     df_reference = pd.read_csv(args.reference)
     logging.info(f"Reference DataFrame shape: {df_reference.shape}")
     logging.info(f"Reference DataFrame columns: {df_reference.columns.tolist()}")
-    
+
     # Cross join the config and reference DataFrames
-    df_in = df_config.merge(df_reference, how='cross')
+    df_in = df_config.merge(df_reference, how="cross")
 
     # TODO: Parallelize the evaluation process for each row in df_in using concurrent.futures or similar libraries
     df_evals = pd.DataFrame()
     for index, row in df_in.iterrows():
+        df_evals = pd.concat(
+            [
+                df_evals,
+                evaluate_response(
+                    row["Model Name"], row["Query"], row["Context"], row["Reference"]
+                ),
+            ],
+            axis=0,
+        )
 
-        #TODO: Rename Model to Model name for clarity
-        df_evals = pd.concat([df_evals, evaluate_response(row['Model'], row['Query'], row['Context'], row['Reference'])], axis=0)
-        
         logging.info(f"Processed row {index + 1}/{len(df_in)}")
 
-    
     # Concatenate the input and evaluations DataFrames
 
-    df_out = pd.concat([df_in.reset_index(drop=True), df_evals.reset_index(drop=True)], axis=1)
+    df_out = pd.concat(
+        [df_in.reset_index(drop=True), df_evals.reset_index(drop=True)], axis=1
+    )
 
     df_out.to_csv(args.output, index=False)
     logging.info(f"Output DataFrame shape: {df_out.shape}")
     logging.info(f"Results saved to {args.output}")
-    logging.info("Evaluation completed successfully.")
+    logging.info("Evaluation completed successfully.")