Remove cuda necessity for SetFit

Author: gabriel-piles

src/trainable_entity_extractor/adapters/extractors/pdf_to_multi_option_extractor/MultiLabelMethod.py

@@ -55,28 +55,19 @@ def load_json(self, file_name: str):
     def remove_model(self):
         shutil.rmtree(join(self.get_path()), ignore_errors=True)
 
-    def get_texts_labels(self, multi_option_data: ExtractionData) -> (list[str], list[list[int]]):
+    def get_texts_labels(self, extraction_data: ExtractionData) -> (list[str], list[list[int]]):
         texts = list()
-        for sample in multi_option_data.samples:
+        for sample in extraction_data.samples:
             texts.append(" ".join([x.text_content.strip() for x in sample.pdf_data.pdf_data_segments]))
 
-        labels = self.get_one_hot_encoding(multi_option_data)
+        labels = self.get_one_hot_encoding(extraction_data)
         return texts, labels
 
-    def predictions_to_options_list(self, predictions) -> list[list[Value]]:
-        return [self.one_prediction_to_option_list(prediction) for prediction in predictions]
-
-    def one_prediction_to_option_list(self, prediction) -> list[Value]:
-        if not self.multi_value:
-            best_score_index = argmax(prediction)
-            return [self.options[best_score_index]] if prediction[best_score_index] > 0.5 else []
-
-        return [Value.from_option(self.options[i]) for i, value in enumerate(prediction) if value > 0.5]
-
-    def get_one_hot_encoding(self, multi_option_data: ExtractionData):
-        options_ids = [option.id for option in self.options]
+    @staticmethod
+    def get_one_hot_encoding(extraction_data: ExtractionData):
+        options_ids = [option.id for option in extraction_data.options]
         one_hot_encoding = list()
-        for sample in multi_option_data.samples:
+        for sample in extraction_data.samples:
             one_hot_encoding.append([0] * len(options_ids))
             for option in sample.labeled_data.values:
                 if option.id not in options_ids:

src/trainable_entity_extractor/adapters/extractors/pdf_to_multi_option_extractor/multi_labels_methods/SetFitEnglishMethod.py

@@ -75,7 +75,7 @@ def train(self, extraction_data: ExtractionData):
 
         model = SetFitModel.from_pretrained(
             self.model_name,
-            labels=[x.label for x in self.options],
+            labels=[x.label for x in extraction_data.options],
             multi_target_strategy="one-vs-rest",
             trust_remote_code=True,
         )
@@ -118,6 +118,8 @@ def predict(self, prediction_samples_data: PredictionSamplesData) -> list[list[V
         if prediction_samples_data.multi_value:
             predictions_proba = model.predict_proba(texts)
             threshold = 0.5
+            if hasattr(predictions_proba, "cpu"):
+                predictions_proba = predictions_proba.cpu().numpy()
             predictions = (predictions_proba > threshold).astype(int)
 
         predictions_values = list()
@@ -137,9 +139,6 @@ def predict(self, prediction_samples_data: PredictionSamplesData) -> list[list[V
         return predictions_values
 
     def can_be_used(self, extraction_data: ExtractionData) -> bool:
-        if not torch.cuda.is_available():
-            return False
-
         if not extraction_data.multi_value:
             return False
 

src/trainable_entity_extractor/adapters/extractors/pdf_to_multi_option_extractor/multi_labels_methods/SingleLabelSetFitEnglishMethod.py

@@ -28,9 +28,6 @@ def gpu_needed(self) -> bool:
         return True
 
     def can_be_used(self, extraction_data: ExtractionData) -> bool:
-        if not torch.cuda.is_available():
-            return False
-
         if extraction_data.multi_value:
             return False
 

src/trainable_entity_extractor/tests/unit_tests/extractors/pdf_to_multi_option_extractor/multi_labels_methods/test_setfit_multilingual.py

@@ -11,6 +11,8 @@
 from trainable_entity_extractor.domain.LabeledData import LabeledData
 from trainable_entity_extractor.domain.Option import Option
 from trainable_entity_extractor.domain.PdfData import PdfData
+from trainable_entity_extractor.domain.PredictionSample import PredictionSample
+from trainable_entity_extractor.domain.PredictionSamplesData import PredictionSamplesData
 from trainable_entity_extractor.domain.TrainingSample import TrainingSample
 from trainable_entity_extractor.adapters.extractors.pdf_to_multi_option_extractor.multi_labels_methods.SetFitMultilingualMethod import (
     SetFitMultilingualMethod,
@@ -25,7 +27,7 @@ class TestSetFitMultilingualMethod(TestCase):
     def tearDown(self):
         shutil.rmtree(join(DATA_PATH, self.TENANT), ignore_errors=True)
 
-    @unittest.SkipTest
+    @unittest.skip("Skipping GPU test in CI/CD")
     def test_train_and_predict(self):
         if not torch.cuda.is_available():
             return
@@ -55,21 +57,19 @@ def test_train_and_predict(self):
         extraction_data = ExtractionData(
             multi_value=True, options=options, samples=samples, extraction_identifier=extraction_identifier
         )
-        setfit_multilingual_method = SetFitMultilingualMethod(extraction_identifier, options, True)
+        setfit_multilingual_method = SetFitMultilingualMethod(extraction_identifier)
 
         try:
             setfit_multilingual_method.train(extraction_data)
         except Exception as e:
             self.fail(f"train() raised {type(e).__name__}")
 
-        prediction_sample_1 = TrainingSample(pdf_data=pdf_data_1)
-        prediction_sample_2 = TrainingSample(pdf_data=pdf_data_2)
-        prediction_sample_4 = TrainingSample(pdf_data=pdf_data_4)
+        prediction_sample_1 = PredictionSample(pdf_data=pdf_data_1)
+        prediction_sample_2 = PredictionSample(pdf_data=pdf_data_2)
+        prediction_sample_4 = PredictionSample(pdf_data=pdf_data_4)
         prediction_samples = [prediction_sample_1, prediction_sample_2, prediction_sample_4]
 
-        prediction_data = ExtractionData(
-            multi_value=True, options=options, samples=prediction_samples, extraction_identifier=extraction_identifier
-        )
+        prediction_data = PredictionSamplesData(multi_value=True, options=options, prediction_samples=prediction_samples)
         predictions = setfit_multilingual_method.predict(prediction_data)
 
         self.assertEqual(3, len(predictions))