refactor: move mappings to data_loaders.py

tasks/data/data_loaders.py

@@ -15,7 +15,7 @@ class DataLoader(ABC):
         pass
 
 class TextDataLoader(DataLoader):
-    def __init__(self, request: TextEvaluationRequest, light: bool = False):
+    def __init__(self, request: TextEvaluationRequest = TextEvaluationRequest(), light: bool = False):
         self.label_mapping = {
             "0_not_relevant": 0,
             "1_not_happening": 1,
@@ -43,3 +43,9 @@ class TextDataLoader(DataLoader):
 
     def get_test_dataset(self):
         return self.dataset["test"]
+
+    def get_label_to_id_mapping(self):
+        return self.label_mapping
+
+    def get_id_to_label_mapping(self):
+        return {v: k for k, v in self.label_mapping.items()}

tasks/models/text_classifiers.py

@@ -8,6 +8,7 @@ from transformers import AutoTokenizer, DataCollatorWithPadding, create_optimize
 import evaluate
 
 from tasks.data.data_loaders import TextDataLoader
+from tasks.utils.evaluation import TextEvaluationRequest
 
 # Define label mappings
 LABEL_TO_ID_MAPPING = {
@@ -23,7 +24,7 @@ LABEL_TO_ID_MAPPING = {
 ID_TO_LABEL_MAPPING = {v: k for k, v in LABEL_TO_ID_MAPPING.items()}
 
 class PredictionModel(ABC):
-    def __init__(self):
+    def __init__(self, data_loader: TextDataLoader = TextDataLoader()):
         self.description = ""
 
     @abstractmethod
@@ -44,12 +45,24 @@ class PredictionModel(ABC):
         pass
 
     @abstractmethod
-    def train(self, dataset):
+    def train(self, dataset) -> None:
+        """
+        Train the model on a given dataset.
+
+        Parameters:
+        -----------
+        dataset:
+            The dataset to train on.
+
+        Returns:
+        --------
+        None
+        """
         pass
 
 
 class BaselineModel(PredictionModel):
-    def __init__(self):
+    def __init__(self, data_loader: TextDataLoader = TextDataLoader()):
         super().__init__()
         self.description = "Random Baseline (with Strategy Pattern, from another module)"
 
@@ -60,10 +73,12 @@ class BaselineModel(PredictionModel):
         pass
 
 class DistilBERTModel(PredictionModel):
-    def __init__(self):
+    def __init__(self, data_loader: TextDataLoader = TextDataLoader()):
         super().__init__()
         self.description = "DistilBERT Model"
         self.model = None
+        self.label_to_id_mapping = data_loader.get_label_to_id_mapping()
+        self.id_to_label_mapping = data_loader.get_id_to_label_mapping()
 
         # tokenizer
         self.tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")
@@ -83,9 +98,10 @@ class DistilBERTModel(PredictionModel):
         logits = outputs.logits
         probabilities = tf.nn.softmax(logits)
         predicted_label = self.model.config.id2label[tf.argmax(probabilities, axis=1).numpy()[0]]
-        return LABEL_TO_ID_MAPPING[predicted_label]
+        return self.label_to_id_mapping[predicted_label]
 
     def train(self, dataset):
+        # Pre-process data
         tokenized_data = self.pre_process_data(dataset)
 
         # Training setup
@@ -99,8 +115,8 @@ class DistilBERTModel(PredictionModel):
         self.model = TFAutoModelForSequenceClassification.from_pretrained(
             "distilbert-base-uncased",
             num_labels=8,
-            id2label=ID_TO_LABEL_MAPPING,
-            label2id=LABEL_TO_ID_MAPPING
+            id2label=self.id_to_label_mapping,
+            label2id=self.label_to_id_mapping
         )
 
         # Convert datasets to tf.data.Dataset format