Fix output path for predictions use case

Author: gabriel-piles

src/trainable_entity_extractor/domain/TrainableEntityExtractorJob.py

@@ -12,3 +12,4 @@ class TrainableEntityExtractorJob(BaseModel):
     options: list[Option] = []
     gpu_needed: bool
     timeout: int
+    output_path: str = ""

src/trainable_entity_extractor/ports/JobExecutor.py

@@ -58,27 +58,6 @@ def recreate_model_folder(extraction_identifier: ExtractionIdentifier) -> None:
         shutil.rmtree(extraction_identifier.get_path(), ignore_errors=True)
         extraction_identifier.get_path().mkdir(parents=True, exist_ok=True)
 
-    def check_and_wait_for_model(self, extraction_identifier: ExtractionIdentifier) -> bool:
-        try:
-            completion_signal_exists = self.model_storage.check_model_completion_signal(extraction_identifier)
-
-            if not completion_signal_exists:
-                self.logger.log(extraction_identifier, "Model completion signal not found, model may still be uploading")
-                return False
-
-            self.logger.log(extraction_identifier, "Model completion signal found, model is ready")
-            model_downloaded = self.model_storage.download_model(extraction_identifier)
-            if model_downloaded:
-                self.logger.log(
-                    extraction_identifier, "Model download failed, checking completion signal", LogSeverity.warning
-                )
-                return True
-
-            return False
-        except Exception as e:
-            self.logger.log(extraction_identifier, f"Error checking model availability: {e}", LogSeverity.error)
-            return False
-
     def is_extractor_cancelled(self, extractor_identifier: ExtractionIdentifier) -> bool:
         try:
             return self.data_retriever.is_extractor_cancelled(extractor_identifier)

src/trainable_entity_extractor/tests/use_cases/test_extractor_text_to_text.py

@@ -87,7 +87,6 @@ def test_predictions_from_source_text_in_labeled_data(self):
         # Create prediction samples
         texts = ["test 0"]
         predictions_samples = [PredictionSample.from_text(text, str(i)) for i, text in enumerate(texts)]
-        predictions_samples[0].segment_selector_texts = []
 
         # Save prediction data
         self.data_retriever.save_prediction_data(extraction_identifier, predictions_samples)

src/trainable_entity_extractor/use_cases/PredictUseCase.py

@@ -14,8 +14,9 @@ def __init__(self, extractors: list[type[ExtractorBase]], logger: Logger):
         self.logger = logger
 
     def predict(self, extractor_job: TrainableEntityExtractorJob, samples: list[PredictionSample]) -> list[Suggestion]:
+        output_path = extractor_job.output_path if extractor_job.output_path else DATA_PATH
         extraction_identifier = ExtractionIdentifier(
-            run_name=extractor_job.run_name, output_path=DATA_PATH, extraction_name=extractor_job.extraction_name
+            run_name=extractor_job.run_name, output_path=output_path, extraction_name=extractor_job.extraction_name
         )
 
         extractor_name = extractor_job.extractor_name