fix

Gapminder · Dec 31, 2024 · c10bf1b · c10bf1b
1 parent d21f585
commit c10bf1b
Showing 1 changed file with 21 additions and 17 deletions.
diff --git a/automation-api/lib/pilot/generate_experiment.py b/automation-api/lib/pilot/generate_experiment.py
@@ -3,14 +3,6 @@
 from typing import Dict, List
 
 import pandas as pd
-
-# List of evaluators to include in experiments
-EVALUATORS: List[str] = [
-    'vertex_ai/gemini-1.5-pro-002',
-    'gpt-4o-2024-11-20', 
-    'anthropic/claude-3-5-sonnet-202401022',
-    'deepseek/deepseek-chat'
-]
 from pandera.errors import SchemaError
 
 from lib.ai_eval_spreadsheet.wrapper import (
@@ -25,18 +17,26 @@
 
 logger = AppSingleton().get_logger()
 
+# List of evaluators to include in experiments
+EVALUATORS: List[str] = [
+    "vertex_ai/gemini-1.5-pro-002",
+    "gpt-4o-2024-11-20",
+    "anthropic/claude-3-5-sonnet-202401022",
+]
+
 
 def filter_included_rows(df: pd.DataFrame) -> pd.DataFrame:
     """
     Filter dataframe to keep only rows where include_in_next_evaluation is True
     and remove the include_in_next_evaluation column
     """
-    if 'include_in_next_evaluation' in df.columns:
-        filtered_df = df[df['include_in_next_evaluation'] == True].copy()
-        filtered_df.drop(columns=['include_in_next_evaluation'], inplace=True)
+    if "include_in_next_evaluation" in df.columns:
+        filtered_df = df[df["include_in_next_evaluation"] is True].copy()
+        filtered_df.drop(columns=["include_in_next_evaluation"], inplace=True)
         return filtered_df
     return df
 
+
 def read_ai_eval_spreadsheet() -> AiEvalData:
     config = read_config()
     authorized_clients = get_service_account_authorized_clients()
@@ -73,19 +73,21 @@ def save_sheets_as_csv() -> Dict[str, str]:
     ai_eval_data = read_ai_eval_spreadsheet()
 
     # Create evaluators DataFrame
-    evaluators_data = {
-        'evaluator': EVALUATORS
-    }
+    evaluators_data = {"evaluator": EVALUATORS}
     evaluators_df = pd.DataFrame(evaluators_data)
 
     # Map of editors to their corresponding sheet names and apply filtering
     editor_map = {
         "evaluators": evaluators_df,
         "questions": filter_included_rows(ai_eval_data.questions.data.df),
         "question_options": filter_included_rows(ai_eval_data.question_options.data.df),
-        "prompt_variations": filter_included_rows(ai_eval_data.prompt_variations.data.df),
+        "prompt_variations": filter_included_rows(
+            ai_eval_data.prompt_variations.data.df
+        ),
         "gen_ai_models": filter_included_rows(ai_eval_data.gen_ai_models.data.df),
-        "gen_ai_model_configs": filter_included_rows(ai_eval_data.gen_ai_model_configs.data.df),
+        "gen_ai_model_configs": filter_included_rows(
+            ai_eval_data.gen_ai_model_configs.data.df
+        ),
         "metrics": filter_included_rows(ai_eval_data.metrics.data.df),
         # no need to export the results
         # "evaluation_results": ai_eval_data.evaluation_results,
@@ -97,7 +99,9 @@ def save_sheets_as_csv() -> Dict[str, str]:
     for sheet_key, df in editor_map.items():
         output_path = os.path.join(base_dir, f"{sheet_key}.csv")
         df.to_csv(output_path, index=False)
-        if sheet_key in sheet_names:  # Only use sheet_names mapping for sheets in the schema
+        if (
+            sheet_key in sheet_names
+        ):  # Only use sheet_names mapping for sheets in the schema
             saved_files[sheet_names[sheet_key]] = output_path
         else:
             saved_files[sheet_key] = output_path