Remove unnecessary classes

pyproject.toml

@@ -1,12 +1,12 @@
 [project]
 name = "trainable-entity-extractor"
-version = "2024.11.1.5003"
+version = "2024.11.1.5004"
 description = "This tool is a trainable text/PDF to entity extractor"
 license = { file = "LICENSE" }
 authors = [{ name = "HURIDOCS" }]
 requires-python = ">= 3.11"
 dependencies = [
-    "pdf-document-layout-analysis @ git+https://github.com/huridocs/pdf-document-layout-analysis@3b990f6dce978feebd59ae0c7de472e5e311387f",
+    "pdf-document-layout-analysis @ git+https://github.com/huridocs/pdf-document-layout-analysis@2f634aa02deb7d1cd53370ec18ee62a325943016",
     "python-Levenshtein==0.25.1",
     "tdda==2.0.9",
     "datasets==2.19.0",

requirements.txt

@@ -1,4 +1,4 @@
-git+https://github.com/huridocs/pdf-document-layout-analysis@949014eb06c1e5010f1fa4e89c8c48ad0ebebea4
+git+https://github.com/huridocs/pdf-document-layout-analysis@2f634aa02deb7d1cd53370ec18ee62a325943016
 python-Levenshtein==0.25.1
 tdda==2.0.9
 datasets==2.19.0

src/trainable_entity_extractor/TrainableEntityExtractor.py

@@ -1,7 +1,7 @@
 from time import time
 
 from trainable_entity_extractor.data.ExtractionIdentifier import ExtractionIdentifier
-from trainable_entity_extractor.data.LogsMessage import Severity
+from trainable_entity_extractor.data.LogSeverity import LogSeverity
 from trainable_entity_extractor.data.PredictionSample import PredictionSample
 from trainable_entity_extractor.data.Suggestion import Suggestion
 from trainable_entity_extractor.extractors.ExtractorBase import ExtractorBase
@@ -50,14 +50,14 @@ def train(self, extraction_data: ExtractionData) -> (bool, str):
             self.extraction_identifier.save_extractor_used(extractor_instance.get_name())
             return extractor_instance.create_model(extraction_data)
 
-        send_logs(self.extraction_identifier, "Error creating extractor", Severity.error)
+        send_logs(self.extraction_identifier, "Error creating extractor", LogSeverity.error)
 
         return False, "Error creating extractor"
 
     def predict(self, prediction_samples: list[PredictionSample]) -> list[Suggestion]:
         extractor_name = self.extraction_identifier.get_extractor_used()
         if not extractor_name:
-            send_logs(self.extraction_identifier, f"No extractor available", Severity.error)
+            send_logs(self.extraction_identifier, f"No extractor available", LogSeverity.error)
             return []
 
         for extractor in self.EXTRACTORS:
@@ -71,5 +71,5 @@ def predict(self, prediction_samples: list[PredictionSample]) -> list[Suggestion
             send_logs(self.extraction_identifier, message)
             return suggestions
 
-        send_logs(self.extraction_identifier, f"No extractor available", Severity.error)
+        send_logs(self.extraction_identifier, f"No extractor available", LogSeverity.error)
         return []

src/trainable_entity_extractor/data/LogSeverity.py

@@ -0,0 +1,6 @@
+from enum import Enum
+
+
+class LogSeverity(str, Enum):
+    error = "error"
+    info = "info"

src/trainable_entity_extractor/extractors/ToTextExtractor.py

@@ -1,7 +1,7 @@
 from trainable_entity_extractor.config import config_logger
 from trainable_entity_extractor.data.ExtractionData import ExtractionData
 from trainable_entity_extractor.data.ExtractionIdentifier import ExtractionIdentifier
-from trainable_entity_extractor.data.LogsMessage import Severity
+from trainable_entity_extractor.data.LogSeverity import LogSeverity
 from trainable_entity_extractor.data.PredictionSample import PredictionSample
 from trainable_entity_extractor.data.Suggestion import Suggestion
 from trainable_entity_extractor.extractors.ExtractorBase import ExtractorBase
@@ -99,7 +99,7 @@ def get_best_method(self, extraction_data: ExtractionData):
                 performance = method_instance.performance(training_set, test_set)
             except Exception as e:
                 message = f"Error checking {method_instance.get_name()}"
-                send_logs(self.extraction_identifier, message, Severity.error, e)
+                send_logs(self.extraction_identifier, message, LogSeverity.error, e)
                 performance = 0
             performance_log += f"{method_instance.get_name()}: {round(performance, 2)}%\n"
             send_logs(self.extraction_identifier, f"Performance {method_instance.get_name()}: {performance}%")

src/trainable_entity_extractor/extractors/pdf_to_multi_option_extractor/PdfToMultiOptionExtractor.py

@@ -4,7 +4,7 @@
 from pathlib import Path
 
 from trainable_entity_extractor.data.ExtractionIdentifier import ExtractionIdentifier
-from trainable_entity_extractor.data.LogsMessage import Severity
+from trainable_entity_extractor.data.LogSeverity import LogSeverity
 from trainable_entity_extractor.data.Option import Option
 from trainable_entity_extractor.data.PredictionSample import PredictionSample
 from trainable_entity_extractor.data.Suggestion import Suggestion
@@ -221,7 +221,7 @@ def get_method_performance(
         try:
             performance = method.get_performance(train_set, test_set)
         except Exception as e:
-            severity = Severity.error if method.REPORT_ERRORS else Severity.info
+            severity = LogSeverity.error if method.REPORT_ERRORS else LogSeverity.info
             send_logs(self.extraction_identifier, f"Error checking {method.get_name()}", severity, e)
             performance = 0
 

src/trainable_entity_extractor/send_logs.py

@@ -2,16 +2,16 @@
 
 from trainable_entity_extractor.config import config_logger
 from trainable_entity_extractor.data.ExtractionIdentifier import ExtractionIdentifier
-from trainable_entity_extractor.data.LogsMessage import Severity
+from trainable_entity_extractor.data.LogSeverity import LogSeverity
 
 
 def send_logs(
     extraction_identifier: ExtractionIdentifier,
     message: str,
-    severity: Severity = Severity.info,
+    severity: LogSeverity = LogSeverity.info,
     exception: Exception = None,
 ):
-    if severity != Severity.error:
+    if severity != LogSeverity.error:
         config_logger.info(message + " for " + extraction_identifier.model_dump_json())
         return