fix data loader

XuehengWang · XuehengWang · commit 2e5004119b11 · 2025-02-20T18:25:29.000Z
diff --git a/tracestorm/cli.py b/tracestorm/cli.py
@@ -133,9 +133,9 @@ def main(
 
         if datasets_config_file is None:
             datasets = []
-            sort = None
+            sort_strategy = None
         else:
-            datasets, sort = load_datasets(datasets_config_file)
+            datasets, sort_strategy = load_datasets(datasets_config_file)
 
         _, result_analyzer = run_load_test(
             trace_generator=trace_generator,
@@ -144,7 +144,7 @@ def main(
             base_url=base_url,
             api_key=api_key,
             datasets=datasets,
-            sort=sort,
+            sort_strategy=sort_strategy,
             seed=seed,
         )
 
diff --git a/tracestorm/core.py b/tracestorm/core.py
@@ -18,7 +18,7 @@ def run_load_test(
     base_url: str,
     api_key: str,
     datasets: List,
-    sort: Optional[str] = None,
+    sort_strategy: Optional[str] = None,
     seed: Optional[int] = None,
 ) -> Tuple[List[Tuple], ResultAnalyzer]:
     """
@@ -31,7 +31,7 @@ def run_load_test(
         base_url: Base URL for API calls
         api_key: API key for authentication
         datasets: List of datasets to generate prompts
-        sort: Sorting strategy for prompts in datasets.
+        sort_strategy: Sorting strategy for prompts in datasets.
         seed: Random seed for sorting.
 
     Returns:
@@ -48,7 +48,7 @@ def run_load_test(
         model_name=model,
         nums=total_requests,
         datasets=datasets,
-        sort=sort,
+        sort_strategy=sort_strategy,
         seed=seed,
     )
     ipc_queue = multiprocessing.Queue()
diff --git a/tracestorm/data_loader.py b/tracestorm/data_loader.py
@@ -44,6 +44,8 @@ def normalize_prompts(row) -> List[str]:
                     "",
                 )
                 prompts.append(prompt)
+            else: # we cannot handle this type
+                continue
     elif isinstance(row, str):  # if the row is already a prompt
         prompts.append(row)
     elif (
@@ -78,109 +80,112 @@ def load_datasets(
     Return:
         (List[Dataset], str): A list of Dataset objects and the sorting strategy.
     """
+    if datasets_config_file is None:
+        logger.error("Customized data loading logic needs to be implemented!")
+        return [], None
+    
     # Load datasets configuration file
-    if datasets_config_file:
-        try:
-            with open(datasets_config_file, "r") as f:
-                datasets_config = json.load(f)
-        except FileNotFoundError:
-            logger.error(
-                f"Configuration file '{datasets_config_file}' not found"
-            )
-            return [], None
-        except Exception as e:
-            logger.error(f"Error reading '{datasets_config_file}': {e}")
-            return [], None
-
-        # Strategy to sort the provided datasets
-        sort_strategy = datasets_config.pop("sort", "random")
-
-        # List to store each Dataset
-        datasets = []
+    try:
+        with open(datasets_config_file, "r") as f:
+            datasets_config = json.load(f)
+    except FileNotFoundError:
+        logger.error(
+            f"Configuration file '{datasets_config_file}' not found"
+        )
+        return [], None
+    except Exception as e:
+        logger.error(f"Error reading '{datasets_config_file}': {e}")
+        return [], None
 
-        for name, config in datasets_config.items():
-            file_name = config.get("file_name")
-            prompt_field = config.get("prompt_field")
+    # Strategy to sort the provided datasets
+    sort_strategy = datasets_config.pop("sort_strategy", "random")
 
-            try:
-                ratio = int(config.get("select_ratio", 1))
-            except ValueError:
-                logger.error(
-                    f"Invalid 'select_ratio' for dataset '{name}', using default 1"
-                )
-                ratio = 1
+    # List to store each Dataset
+    datasets = []
 
-            if not file_name or not prompt_field:
-                logger.error(
-                    f"Missing required 'file_name' or 'prompt_field' for dataset '{name}'"
-                )
-                continue
+    for name, config in datasets_config.items():
+        file_name = config.get("file_name")
+        prompt_field = config.get("prompt_field")
 
-            file_path = (
-                os.path.abspath(file_name)
-                if os.path.exists(file_name)
-                else os.path.join(DEFAULT_DATASET_FOLDER, file_name)
+        try:
+            ratio = int(config.get("select_ratio", 1))
+        except ValueError:
+            logger.error(
+                f"Invalid 'select_ratio' for dataset '{name}', using default 1"
             )
+            ratio = 1
 
-            # Load dataset from local files
-            if os.path.exists(file_path):
-                prompts = []
-                # CSV files
-                if file_name.endswith(".csv"):
-                    data = pd.read_csv(file_path)
+        if not file_name or not prompt_field:
+            logger.error(
+                f"Missing required 'file_name' or 'prompt_field' for dataset '{name}'"
+            )
+            continue
+        
+        os.makedirs(DEFAULT_DATASET_FOLDER, exist_ok=True)
+        
+        file_path = (
+            os.path.abspath(file_name)
+            if os.path.exists(file_name)
+            else os.path.join(DEFAULT_DATASET_FOLDER, file_name)
+        )
 
-                    if prompt_field not in set(data.columns):
+        # Load dataset from local files
+        if os.path.exists(file_path):
+            prompts = []
+            # CSV files
+            if file_name.endswith(".csv"):
+                data = pd.read_csv(file_path)
+
+                if prompt_field not in set(data.columns):
+                    logger.error(
+                        f"Field '{prompt_field}' not found in '{file_path}'."
+                    )
+                    continue
+                prompts = data[prompt_field].dropna().astype(str).tolist()
+            # JSON files
+            elif file_name.endswith(".json"):
+                with open(file_path, "r") as f:
+                    data = json.load(f)
+
+                if isinstance(data, dict):
+                    prompts = data.get(prompt_field, [])
+                    if not isinstance(prompts, list):
                         logger.error(
-                            f"Field '{prompt_field}' not found in '{file_path}'."
+                            f"Field '{prompt_field}' in '{file_path}' is not a list."
                         )
                         continue
-                    prompts = data[prompt_field].dropna().astype(str).tolist()
-                # JSON files
-                elif file_name.endswith(".json"):
-                    with open(file_path, "r") as f:
-                        data = json.load(f)
-
-                    if isinstance(data, dict):
-                        prompts = data.get(prompt_field, [])
-                        if not isinstance(prompts, list):
-                            logger.error(
-                                f"Field '{prompt_field}' in '{file_path}' is not a list."
-                            )
-                            continue
-                else:
-                    logger.error(f"Unsupported file format for '{file_name}'")
-                    continue
             else:
-                try:
-                    if file_name.endswith(".csv"):  # CSV format
-                        data = pd.read_csv(file_name)
-
-                        if prompt_field not in set(data.columns):
-                            logger.error(
-                                f"Field '{prompt_field}' not found in '{file_name}'."
-                            )
-                            continue
-                        prompts = (
-                            data[prompt_field].dropna().astype(str).tolist()
+                logger.error(f"Unsupported file format for '{file_name}'")
+                continue
+        else:
+            try:
+                if file_name.endswith(".csv"):  # CSV format
+                    data = pd.read_csv(file_name)
+
+                    if prompt_field not in set(data.columns):
+                        logger.error(
+                            f"Field '{prompt_field}' not found in '{file_name}'."
                         )
-                    else:  # use datasets to load
-                        data = load_dataset(file_name)["train"]
-                        prompts = []
-                        for row in data[prompt_field]:
-                            prompts.extend(normalize_prompts(row))
-                except Exception as e:
-                    logger.error(f"Failed to load '{file_name}': {e}")
-
-            # Add the dataset information (file name, a list of prompts, select ratio among all datasets, total number of prompts)
-            dataset_obj = Dataset(file_name, prompts, ratio, len(prompts))
-            datasets.append(dataset_obj)
-
-            logger.info(
-                f"loaded {file_name} with {len(prompts)} prompts, selection ratio = {ratio}"
-            )
+                        continue
+                    prompts = (
+                        data[prompt_field].dropna().astype(str).tolist()
+                    )
+                else:  # use datasets to load
+                    data = load_dataset(file_name)["train"]
+                    prompts = []
+                    for row in data[prompt_field]:
+                        prompts.extend(normalize_prompts(row))
+            except Exception as e:
+                logger.error(f"Failed to load '{file_name}': {e}")
+
+        # Add the dataset information (file name, a list of prompts, select ratio among all datasets, total number of prompts)
+        dataset_obj = Dataset(file_name, prompts, ratio, len(prompts))
+        datasets.append(dataset_obj)
+
+        logger.info(
+            f"loaded {file_name} with {len(prompts)} prompts, selection ratio = {ratio}"
+        )
 
-        return datasets, sort_strategy
+    return datasets, sort_strategy
 
-    else:
-        logger.error("Customized data loading logic needs to be implemented!")
-        return [], None
+        
diff --git a/tracestorm/request_generator.py b/tracestorm/request_generator.py
@@ -13,7 +13,7 @@ def generate_request(
     nums: int,
     messages: str = DEFAULT_MESSAGES,
     datasets: List[Dataset] = [],
-    sort: str = "random",
+    sort_strategy: str = "random",
     seed: int = None,
 ) -> List[Dict[str, Any]]:
     # generate default requests without datasets
@@ -58,14 +58,14 @@ def generate_request(
             )
 
         # 1. Randomly sort the requests
-        if sort == "random":
+        if sort_strategy == "random":
             if seed is not None:
                 random.seed(seed)
             random.shuffle(dataset_samples)
-        elif sort == "original":  # 2. original order
+        elif sort_strategy == "original":  # 2. original order
             dataset_samples.sort(key=lambda x: x[0])
         else:
-            raise ValueError(f"Unknown sort strategy: {sort}")
+            raise ValueError(f"Unknown sorting strategy: {sort_strategy}")
 
         # Extract the prompts from the list
         requests = [
