Add metadata to trainable extraction jobs

Author: gabriel-piles

src/trainable_entity_extractor/domain/TrainableEntityExtractorJob.py

@@ -13,6 +13,7 @@ class TrainableEntityExtractorJob(BaseModel):
     gpu_needed: bool
     timeout: int
     output_path: str = ""
+    metadata: dict[str, str] = dict()
 
     def set_extractors_path(self, path: str) -> "TrainableEntityExtractorJob":
         self.output_path = path

src/trainable_entity_extractor/ports/ExtractorBase.py

@@ -110,6 +110,9 @@ def get_distributed_jobs(self, extraction_data: ExtractionData) -> list[Trainabl
                 timeout=getattr(method_instance, "timeout", 3600),
                 options=extraction_data.options if extraction_data.options else [],
                 multi_value=extraction_data.multi_value if extraction_data.multi_value else False,
+                metadata=(
+                    extraction_data.extraction_identifier.metadata if extraction_data.extraction_identifier.metadata else {}
+                ),
             )
             jobs.append(job)
 

src/trainable_entity_extractor/ports/ModelStorage.py

@@ -58,7 +58,7 @@ def serialize_job_to_dict(job: TrainableEntityExtractorJob) -> dict:
             "options": [option.model_dump() for option in job.options],
             "gpu_needed": job.gpu_needed,
             "timeout": job.timeout,
-            "metadata": {},
+            "metadata": job.metadata,
         }
 
     @staticmethod
@@ -74,6 +74,7 @@ def deserialize_job_from_dict(job_data: dict) -> TrainableEntityExtractorJob:
         options = [Option(**option_data) for option_data in options_data]
         gpu_needed = job_data.get("gpu_needed", False)
         timeout = job_data.get("timeout", 3600)
+        metadata = job_data.get("metadata", {})
 
         if version != "1.0":
             pass
@@ -87,4 +88,5 @@ def deserialize_job_from_dict(job_data: dict) -> TrainableEntityExtractorJob:
             options=options,
             gpu_needed=gpu_needed,
             timeout=timeout,
+            metadata=metadata,
         )

src/trainable_entity_extractor/tests/use_cases/test_get_distributed_tasks.py

@@ -110,23 +110,6 @@ def test_get_distributed_jobs_no_compatible_extractor(self):
 
         self.assertEqual(len(jobs), 0)
 
-    def test_get_distributed_jobs_multiple_extractors_returns_last_compatible(self):
-        text_extraction_data = self.create_text_extraction_data()
-        multi_option_extraction_data = self.create_multi_option_extraction_data()
-
-        self.train_use_case = TrainUseCase(extractors=[TextToTextExtractor, TextToMultiOptionExtractor], logger=self.logger)
-
-        text_jobs = self.train_use_case.get_jobs(text_extraction_data)
-        multi_option_jobs = self.train_use_case.get_jobs(multi_option_extraction_data)
-
-        self.assertGreater(len(text_jobs), 0)
-        for job in text_jobs:
-            self.assertEqual(job.extractor_name, "TextToTextExtractor")
-
-        self.assertGreater(len(multi_option_jobs), 0)
-        for job in multi_option_jobs:
-            self.assertEqual(job.extractor_name, "TextToMultiOptionExtractor")
-
     def test_get_distributed_jobs_with_incompatible_data_for_text_extractor(self):
         extraction_data = ExtractionData(
             samples=[],