Refacto/Updated format of Harness

langtest/augmentation/__init__.py

@@ -113,7 +113,8 @@ def fix(
         Returns:
             List[Dict[str, Any]]: A list of augmented data samples.
         """
-        if "." not in training_data["data_source"]:
+
+        if "source" in training_data and training_data["source"] == "huggingface":
             self.df = HuggingFaceDataset(training_data["data_source"], self.task)
             data = self.df.load_data(
                 feature_column=training_data.get("feature_column", "text"),
@@ -122,7 +123,7 @@ def fix(
                 subset=training_data.get("subset", None),
             )
         else:
-            self.df = DataFactory(training_data["data_source"], self.task)
+            self.df = DataFactory(training_data, self.task)
             data = self.df.load()
         TestFactory.is_augment = True
         supported_tests = TestFactory.test_scenarios()
@@ -345,7 +346,7 @@ def fix(
         Returns:
             bool: Returns True upon successful completion of the method.
         """
-        df = DataFactory(training_data["data_source"], self.__task)
+        df = DataFactory(training_data, self.__task)
         data = df.load()
         new_data = []
         self.__search_results = self.search_sample_results(data)

langtest/datahandler/datasource.py

@@ -94,14 +94,21 @@ class DataFactory:
     correct Dataset type based on the file extension.
     """
 
-    def __init__(self, file_path: str, task: str, **kwargs) -> None:
+    def __init__(self, file_path: dict, task: str, **kwargs) -> None:
         """Initializes DataFactory object.
 
         Args:
-            file_path (str): Path to the dataset.
+            file_path (dict): Dictionary containing 'data_source' key with the path to the dataset.
             task (str): Task to be evaluated.
         """
-        self._file_path = file_path
+        if not isinstance(file_path, dict):
+            raise ValueError("'file_path' must be a dictionary.")
+
+        if "data_source" not in file_path:
+            raise ValueError(
+                "The 'data_source' key must be provided in the 'file_path' dictionary."
+            )
+        self._file_path = file_path.get("data_source")
         self._class_map = {
             cls.__name__.replace("Dataset", "").lower(): cls
             for cls in _IDataset.__subclasses__()

langtest/langtest.py

@@ -3,7 +3,7 @@
 import os
 import pickle
 from collections import defaultdict
-from typing import Any, Dict, List, Optional, Union
+from typing import Dict, List, Optional, Union
 
 import pandas as pd
 import yaml
@@ -91,18 +91,18 @@ class Harness:
     def __init__(
         self,
         task: str,
-        model: Optional[Union[str, Any]] = None,
-        hub: Optional[str] = None,
-        data: Optional[Union[str, dict]] = None,
+        model: Optional[Union[list, dict]] = None,
+        data: Optional[dict] = None,
         config: Optional[Union[str, dict]] = None,
     ):
         """Initialize the Harness object.
 
         Args:
             task (str, optional): Task for which the model is to be evaluated.
-            model (str | ModelFactory): ModelFactory object or path to the model to be evaluated.
-            hub (str, optional): model hub to load from the path. Required if path is passed as 'model'.
-            data (str, optional): Path to the data to be used for evaluation.
+            model (list | dict, optional): Specifies the model to be evaluated.
+                If provided as a list, each element should be a dictionary with 'model' and 'hub' keys.
+                If provided as a dictionary, it must contain 'model' and 'hub' keys when specifying a path.
+            data (dict, optional): The data to be used for evaluation.
             config (str | dict, optional): Configuration for the tests to be performed.
 
         Raises:
@@ -111,8 +111,27 @@ def __init__(
         super().__init__()
 
         self.is_default = False
-        self._actual_model = model
-        self.hub = hub
+
+        if isinstance(model, list):
+            for item in model:
+                if not isinstance(item, dict):
+                    raise ValueError("Each item in the list must be a dictionary")
+                if "model" not in item or "hub" not in item:
+                    raise ValueError(
+                        "Each dictionary in the list must have 'model' and 'hub' keys"
+                    )
+        elif isinstance(model, dict):
+            if "model" not in model or "hub" not in model:
+                raise ValueError("The dictionary must have 'model' and 'hub' keys")
+        else:
+            raise ValueError("Invalid 'model' parameter type")
+
+        if isinstance(model, dict):
+            hub, model = model["hub"], model["model"]
+            self.hub = hub
+            self._actual_model = model
+        else:
+            hub = None
 
         if task not in self.SUPPORTED_TASKS:
             raise ValueError(
@@ -133,7 +152,7 @@ def __init__(
 
         if data is None and (task, model, hub) in self.DEFAULTS_DATASET:
             data_path = os.path.join("data", self.DEFAULTS_DATASET[(task, model, hub)])
-            data = resource_filename("langtest", data_path)
+            data = {"data_source": resource_filename("langtest", data_path)}
             self.data = DataFactory(data, task=self.task).load()
             if model == "textcat_imdb":
                 model = resource_filename("langtest", "data/textcat_imdb")
@@ -142,61 +161,60 @@ def __init__(
 
         elif (
             isinstance(data, dict)
-            and hub in self.SUPPORTED_HUBS_HF_DATASET_CLASSIFICATION
-            and task == "text-classification"
+            and "source" in data
+            and data["source"] == "huggingface"
         ):
-            self.data = (
-                HuggingFaceDataset(data["name"], task=task).load_data(
+            if (
+                task == "text-classification"
+                and hub in self.SUPPORTED_HUBS_HF_DATASET_CLASSIFICATION
+            ):
+                self.data = HuggingFaceDataset(data["data_source"], task=task).load_data(
                     feature_column=data.get("feature_column", "text"),
                     target_column=data.get("target_column", "label"),
                     split=data.get("split", "test"),
                     subset=data.get("subset", None),
                 )
-                if data is not None
-                else None
-            )
 
-            if hub == "spacy" and (model == "textcat_imdb" or model is None):
-                if model is None:
-                    logging.warning(
-                        "Using the default 'textcat_imdb' model for Spacy hub. Please provide a custom model path if desired."
-                    )
-                model = resource_filename("langtest", "data/textcat_imdb")
+                if hub == "spacy" and (model == "textcat_imdb" or model is None):
+                    if model is None:
+                        logging.warning(
+                            "Using the default 'textcat_imdb' model for Spacy hub. Please provide a custom model path if desired."
+                        )
+                    model = resource_filename("langtest", "data/textcat_imdb")
 
-        elif (
-            isinstance(data, dict)
-            and hub in self.SUPPORTED_HUBS_HF_DATASET_NER
-            and task == "ner"
-        ):
-            self.data = HuggingFaceDataset(data["name"], task=task).load_data(
-                feature_column=data.get("feature_column", "tokens"),
-                target_column=data.get("target_column", "ner_tags"),
-                split=data.get("split", "test"),
-                subset=data.get("subset", None),
-            )
+            elif task == "ner" and hub in self.SUPPORTED_HUBS_HF_DATASET_NER:
+                self.data = HuggingFaceDataset(data["data_source"], task=task).load_data(
+                    feature_column=data.get("feature_column", "tokens"),
+                    target_column=data.get("target_column", "ner_tags"),
+                    split=data.get("split", "test"),
+                    subset=data.get("subset", None),
+                )
 
-        elif (
-            isinstance(data, dict)
-            and hub in self.SUPPORTED_HUBS_HF_DATASET_SUMMARIZATION
-            and task == "summarization"
-        ):
-            self.data = HuggingFaceDataset(data["name"], task=task).load_data(
-                feature_column=data.get("feature_column", "document"),
-                target_column=data.get("target_column", "summary"),
-                split=data.get("split", "test"),
-                subset=data.get("subset", None),
-            )
+            elif (
+                task == "summarization"
+                and hub in self.SUPPORTED_HUBS_HF_DATASET_SUMMARIZATION
+            ):
+                self.data = HuggingFaceDataset(data["data_source"], task=task).load_data(
+                    feature_column=data.get("feature_column", "document"),
+                    target_column=data.get("target_column", "summary"),
+                    split=data.get("split", "test"),
+                    subset=data.get("subset", None),
+                )
 
         elif data is None and (task, model, hub) not in self.DEFAULTS_DATASET.keys():
             raise ValueError(
                 "You haven't specified any value for the parameter 'data' and the configuration you "
                 "passed is not among the default ones. You need to either specify the parameter 'data' "
                 "or use a default configuration."
             )
-        elif isinstance(data, list):
-            self.data = data
+        elif isinstance(data["data_source"], list):
+            self.data = data["data_source"]
         else:
-            self.file_path = data
+            if "data_source" not in data:
+                raise ValueError(
+                    "The 'data_source' key must be provided in the 'data' parameter."
+                )
+            self.file_path = data["data_source"]
             self.data = (
                 DataFactory(data, task=self.task).load() if data is not None else None
             )
@@ -221,13 +239,19 @@ def __init__(
                 path=model, task=task, hub=hub, **self._config.get("model_parameters", {})
             )
 
-        elif type(model) == dict:
+        elif isinstance(model, list):
             model_dict = {}
-            for k, v in model.items():
-                model_dict[k] = ModelFactory.load_model(
-                    task=task, path=k, hub=v, **self._config.get("model_parameters", {})
+            for i in model:
+                model = i["model"]
+                hub = i["hub"]
+
+                model_dict[model] = ModelFactory.load_model(
+                    path=model,
+                    task=task,
+                    hub=hub,
+                    **self._config.get("model_parameters", {}),
                 )
-            self.model = model_dict
+                self.model = model_dict
 
         else:
             self.model = ModelFactory(
@@ -241,7 +265,7 @@ def __init__(
         print("Test Configuration : \n", formatted_config)
 
         global GLOBAL_MODEL
-        if not isinstance(model, dict):
+        if not isinstance(model, list):
             GLOBAL_MODEL = self.model
 
         self._testcases = None
@@ -290,13 +314,17 @@ def configure(self, config: Union[str, dict]) -> dict:
                     **self._config.get("model_parameters", {}),
                 )
 
-            elif isinstance(model, dict):
+            elif isinstance(model, list):
                 model_dict = {}
-                for k, v in model.items():
-                    model_dict[k] = ModelFactory.load_model(
+
+                for i in model:
+                    model = i["model"]
+                    hub = i["hub"]
+
+                    model_dict[model] = ModelFactory.load_model(
+                        path=model,
                         task=task,
-                        path=k,
-                        hub=v,
+                        hub=hub,
                         **self._config.get("model_parameters", {}),
                     )
                 self.model = model_dict
@@ -1017,9 +1045,8 @@ def load(
 
         harness = Harness(
             task=task,
-            model=model,
-            data=data,
-            hub=hub,
+            model={"model": model, "hub": hub},
+            data={"data_source": data},
             config=os.path.join(save_dir, "config.yaml"),
         )
         harness.generate()
@@ -1050,7 +1077,9 @@ def import_edited_testcases(self, input_path: str, **kwargs):
             if sample.category not in ["robustness", "bias"]
         ]
 
-        self._testcases = DataFactory(input_path, task=self.task, is_import=True).load()
+        self._testcases = DataFactory(
+            {"data_source": input_path}, task=self.task, is_import=True
+        ).load()
         self._testcases.extend(temp_testcases)
 
         return self

langtest/pipelines/transformers/ner_pipeline.py

@@ -84,8 +84,12 @@ def setup(self):
         self.hub = "huggingface"
         self.output_dir = "checkpoints/"
 
-        self.train_datasource = DataFactory(file_path=self.train_data, task=self.task)
-        self.eval_datasource = DataFactory(file_path=self.eval_data, task=self.task)
+        self.train_datasource = DataFactory(
+            file_path={"data_source": self.train_data}, task=self.task
+        )
+        self.eval_datasource = DataFactory(
+            file_path={"data_source": self.eval_data}, task=self.task
+        )
 
         self.next(self.train)
 
@@ -153,9 +157,8 @@ def test(self):
         """Performs the testing procedure of the model on a set of tests using langtest"""
         self.harness = Harness(
             task=self.task,
-            model=self.output_dir,
-            hub=self.hub,
-            data=self.train_data,
+            model={"model": self.output_dir, "hub": self.hub},
+            data={"data_source": self.train_data},
         )
         if self.config:
             self.harness.configure(self.config)
@@ -184,7 +187,7 @@ def augment(self):
     def retrain(self):
         """Performs the training procedure using the augmented data created by langtest"""
         self.augmented_train_datasource = DataFactory(
-            file_path=self.path_augmented_file, task=self.task
+            file_path={"data_source": self.path_augmented_file}, task=self.task
         )
         samples = self.augmented_train_datasource.load_raw()