Added filename

Author: alex-aipolabs

backend/evals/evaluation_pipeline.py

@@ -12,6 +12,7 @@
 logger = logging.getLogger(__name__)
 
 DEFAULT_DATASET_ARTIFACT = "synthetic_intent_dataset"
+DEFAULT_DATASET_FILENAME = "synthetic_intents.csv"
 DEFAULT_EVALUATION_MODEL = "dual-encoder-text-embedding-1024"
 
 
@@ -60,21 +61,26 @@ def __init__(
             api_key=search_api_key,
         )
 
-    def _load_dataset_from_wandb(self, artifact_name: str) -> pd.DataFrame:
+    def _load_dataset_from_wandb(self, artifact_name: str, dataset_filename: str) -> pd.DataFrame:
         """
         Load a dataset from a W&B artifact.
 
         Args:
             artifact_name: Name of the W&B artifact
-
+            dataset_filename: Filename to save the dataset to
         Returns:
             DataFrame containing the dataset
         """
         artifact = wandb.use_artifact(f"{artifact_name}:latest")
         artifact_dir = artifact.download()
-        return pd.read_csv(os.path.join(artifact_dir, "temp_dataset.csv"))
+        return pd.read_csv(os.path.join(artifact_dir, dataset_filename))
 
-    def _generate(self, dataset_artifact: str, generation_limit: int | None = None) -> pd.DataFrame:
+    def _generate(
+        self,
+        dataset_artifact: str,
+        dataset_filename: str,
+        generation_limit: int | None = None,
+    ) -> pd.DataFrame:
         """
         Generate synthetic intents.
 
@@ -88,6 +94,7 @@ def _generate(self, dataset_artifact: str, generation_limit: int | None = None)
         logger.info("Generating synthetic intents...")
         df = self.generator.generate(
             dataset_artifact=dataset_artifact,
+            dataset_filename=dataset_filename,
             limit=generation_limit,
         )
 
@@ -97,6 +104,7 @@ def _generate(self, dataset_artifact: str, generation_limit: int | None = None)
     def _evaluate(
         self,
         dataset_artifact: str,
+        dataset_filename: str,
         evaluation_samples: int | None = None,
         df: pd.DataFrame | None = None,
     ) -> dict:
@@ -113,7 +121,7 @@ def _evaluate(
         """
         if df is None:
             logger.info(f"Loading dataset from artifact: {dataset_artifact}")
-            df = self._load_dataset_from_wandb(dataset_artifact)
+            df = self._load_dataset_from_wandb(dataset_artifact, dataset_filename)
 
         # Evaluate search performance
         logger.info("Evaluating search performance...")
@@ -138,9 +146,10 @@ def _evaluate(
 
     def run(
         self,
+        dataset_artifact: str,
+        dataset_filename: str,
         generate_data: bool = False,
         evaluate_data: bool = True,
-        dataset_artifact: str = DEFAULT_DATASET_ARTIFACT,
         generation_limit: int | None = None,
         evaluation_samples: int | None = None,
     ) -> None:
@@ -169,6 +178,7 @@ def run(
                 "evaluation_model": DEFAULT_EVALUATION_MODEL,
                 "evaluation_samples": evaluation_samples,
                 "dataset_artifact": dataset_artifact,
+                "dataset_filename": dataset_filename,
             },
         )
 
@@ -177,12 +187,14 @@ def run(
             if generate_data:
                 df = self._generate(
                     dataset_artifact=dataset_artifact,
+                    dataset_filename=dataset_filename,
                     generation_limit=generation_limit,
                 )
 
             if evaluate_data:
                 self._evaluate(
                     dataset_artifact=dataset_artifact,
+                    dataset_filename=dataset_filename,
                     evaluation_samples=evaluation_samples,
                     df=df,
                 )
@@ -199,15 +211,26 @@ def run(
     required=True,
 )
 @click.option(
-    "--dataset",
+    "--dataset-artifact",
     default=DEFAULT_DATASET_ARTIFACT,
     help="Name of the W&B dataset artifact to use",
     show_default=True,
 )
+@click.option(
+    "--dataset-filename",
+    default=DEFAULT_DATASET_FILENAME,
+    type=str,
+    help="Filename to save the generated dataset to",
+    show_default=True,
+)
 @click.option("--generation-limit", type=int, help="Limit number of samples to generate")
 @click.option("--evaluation-samples", type=int, help="Limit number of samples to evaluate")
 def main(
-    mode: str, dataset: str, generation_limit: int | None, evaluation_samples: int | None
+    mode: str,
+    dataset_artifact: str,
+    generation_limit: int | None,
+    evaluation_samples: int | None,
+    dataset_filename: str,
 ) -> None:
     """Main entry point for the evaluation pipeline."""
     # Get API keys from environment
@@ -235,9 +258,10 @@ def main(
 
     # Run pipeline
     pipeline.run(
+        dataset_artifact=dataset_artifact,
+        dataset_filename=dataset_filename,
         generate_data=generate_data,
         evaluate_data=evaluate_data,
-        dataset_artifact=dataset,
         generation_limit=generation_limit,
         evaluation_samples=evaluation_samples,
     )

backend/evals/synthetic_intent_generator.py

@@ -1,6 +1,3 @@
-import os
-import tempfile
-
 import openai
 import pandas as pd
 import wandb
@@ -112,7 +109,7 @@ def _log_dataset_stats(self, df: pd.DataFrame) -> None:
             }
         )
 
-    def _save_to_wandb(self, df: pd.DataFrame, dataset_artifact: str) -> str:
+    def _save_to_wandb(self, df: pd.DataFrame, dataset_artifact: str, dataset_filename: str) -> str:
         """
         Save the dataset as a wandb artifact.
 
@@ -142,27 +139,19 @@ def _save_to_wandb(self, df: pd.DataFrame, dataset_artifact: str) -> str:
             },
         )
 
-        # Use tempfile to create and manage a temporary file
-        with tempfile.NamedTemporaryFile(suffix=".csv", delete=False) as temp_file:
-            temp_filename = temp_file.name
-
-        try:
-            # Write dataframe to the temporary file
-            df.to_csv(temp_filename, index=False)
-            # Add the file to the artifact
-            artifact.add_file(temp_filename)
-            # Log the artifact
-            wandb.log_artifact(artifact)
+        # Write dataframe to the temporary file
+        df.to_csv(dataset_filename, index=False)
+        # Add the file to the artifact
+        artifact.add_file(dataset_filename)
+        # Log the artifact
+        wandb.log_artifact(artifact)
 
-            return artifact.name
-        finally:
-            # Ensure temp file is removed even if any operation fails
-            if os.path.exists(temp_filename):
-                os.unlink(temp_filename)
+        return artifact.name
 
     def generate(
         self,
         dataset_artifact: str,
+        dataset_filename: str,
         limit: int | None = None,
     ) -> pd.DataFrame:
         """
@@ -191,7 +180,7 @@ def generate(
 
         # Log and save
         self._log_dataset_stats(df)
-        artifact_name = self._save_to_wandb(df, dataset_artifact)
+        artifact_name = self._save_to_wandb(df, dataset_artifact, dataset_filename)
 
         print(f"Dataset saved as W&B artifact: {artifact_name}")
         return df