refactor: train and save models only if needed

main.py

@@ -15,12 +15,13 @@ async def main():
     data_loader = TextDataLoader(text_request, light=light_dataset)
 
     # define model
-    model = ModelFactory.create_model("distilbert")
+    model = ModelFactory.create_model("baseline")
 
     # train model
     train_dataset = data_loader.get_train_dataset()
-    model.train(train_dataset)
-    model.save()
+    if model.model is None:
+        model.train(train_dataset)
+        model.save()
 
     # Call the evaluate_text function
     results = await evaluate_text(request = text_request,

tasks/models/text_classifiers.py

@@ -13,6 +13,7 @@ from tasks.data.data_loaders import TextDataLoader
 class PredictionModel(ABC):
     def __init__(self, data_loader: TextDataLoader = TextDataLoader()):
         self.description = ""
+        self.model = None
 
     @abstractmethod
     def predict(self, quote: str) -> int:
@@ -66,7 +67,6 @@ class DistilBERTModel(PredictionModel):
     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()