some improvements

src/image_classification/image_classification_trainer.py

@@ -52,6 +52,9 @@ class ImageClassificationTrainer(AbstractTrainer):
 
             self.__train_model(images, parameters)
 
+            if(self.get_status().is_training_aborted()):
+                return
+
             self.get_status().update_status(100, "Training completed")
 
         except Exception as e:

src/main.py

@@ -75,7 +75,7 @@ async def get_task_status(token_data: dict = Depends(verify_token)):
         "status": status.get_status().value
     }
 
-@app.get("/stop_training")
+@app.put("/stop_training")
 async def stop_task(token_data: dict = Depends(verify_token)):
     """ Stop the currently running training (if any). """
     try: 

src/text_classification/text_classification_trainer.py

@@ -44,6 +44,9 @@ class TextClassificationTrainer(AbstractTrainer):
 
             self.__train_model(tokenized_dataset, labels, label2id, id2label, parameters)
 
+            if(self.get_status().is_training_aborted()):
+                return
+
             self.get_status().update_status(100, "Training completed")
 
         except Exception as e:
@@ -66,6 +69,18 @@ class TextClassificationTrainer(AbstractTrainer):
         dataset = load_dataset('csv', data_files=parameters.get_training_csv_file_path(), delimiter=parameters.get_training_csv_limiter())
 
         dataset = dataset["train"]
+
+        # Extract the labels
+        #labels = tokenized_dataset['train'].unique('target')
+        labels = dataset.unique('target')
+        label2id, id2label = dict(), dict()
+        for i, label in enumerate(labels):
+            label2id[label] = i
+            id2label[i] = label
+
+        logger.info(id2label)
+        
+
         dataset = dataset.train_test_split(test_size=0.2)
 
         logger.info(dataset)
@@ -79,15 +94,6 @@ class TextClassificationTrainer(AbstractTrainer):
 
         tokenized_dataset = dataset.map(preprocess_function, batched=True)
 
-        # Extract the labels
-        labels = tokenized_dataset['train'].unique('target')
-        label2id, id2label = dict(), dict()
-        for i, label in enumerate(labels):
-            label2id[label] = i
-            id2label[i] = label
-
-        logger.info(id2label)
-
         # Rename the Target column to labels and remove unnecessary columns
         tokenized_dataset = tokenized_dataset.rename_column('target', 'labels')