Fix Gliner slowness

Author: gabriel-piles

src/trainable_entity_extractor/adapters/extractors/GlinerDateExtractor.py

@@ -1,10 +1,12 @@
 import json
-
 from dateparser.search import search_dates
-from gliner import GLiNER
 
 
 class GlinerDateExtractor:
+
+    def __init__(self, model):
+        self.model = model
+
     @staticmethod
     def find_unique_entity_dicts(entities: list[dict]) -> list[dict]:
         dicts_without_score = [{k: v for k, v in d.items() if k != "score"} for d in entities]
@@ -25,8 +27,6 @@ def remove_overlapping_entities(entities):
         return result
 
     def extract_dates(self, text: str):
-        gliner_model = GLiNER.from_pretrained("urchade/gliner_multi-v2.1")
-
         words = text.split()
 
         entities = []
@@ -37,7 +37,7 @@ def extract_dates(self, text: str):
         for i in range(0, len(words), slide_size):
             window_words = words[i : i + window_size]
             window_text = " ".join(window_words)
-            window_entities = gliner_model.predict_entities(window_text, ["date"])
+            window_entities = self.model.predict_entities(window_text, ["date"])
 
             for entity in window_entities:
                 entity["start"] += last_slide_end_index

src/trainable_entity_extractor/adapters/extractors/pdf_to_text_extractor/methods/GlinerFirstDateMethod.py

@@ -1,26 +1,54 @@
 import re
 
+from gliner import GLiNER
+
+from trainable_entity_extractor.adapters.extractors.ToTextExtractorMethod import ToTextExtractorMethod
+from trainable_entity_extractor.domain.ExtractionData import ExtractionData
 from trainable_entity_extractor.domain.PdfDataSegment import PdfDataSegment
-from trainable_entity_extractor.adapters.extractors.pdf_to_text_extractor.methods.FirstDateMethod import FirstDateMethod
 from trainable_entity_extractor.adapters.extractors.text_to_text_extractor.methods.GlinerDateParserMethod import (
     GlinerDateParserMethod,
 )
+from trainable_entity_extractor.domain.PredictionSamplesData import PredictionSamplesData
+
+
+class GlinerFirstDateMethod(ToTextExtractorMethod):
+    def train(self, extraction_data: ExtractionData):
+        languages = [x.labeled_data.language_iso for x in extraction_data.samples]
+        self.save_json("languages.json", list(set(languages)))
+
+    def predict(self, prediction_samples_data: PredictionSamplesData) -> list[str]:
+        gliner_model = GLiNER.from_pretrained("urchade/gliner_multi-v2.1")
+        predictions_samples = prediction_samples_data.prediction_samples
+        predictions = [""] * len(predictions_samples)
+        languages = self.load_json("languages.json")
+        for index, prediction_sample in enumerate(predictions_samples):
+            segments = prediction_sample.pdf_data.pdf_data_segments
+
+            if predictions[index] or not prediction_sample.pdf_data or not segments:
+                continue
+
+            predictions[index] = self.get_date_from_segments(gliner_model, segments, languages)
+
+        return predictions
 
+    @staticmethod
+    def loop_segments(segments: list[PdfDataSegment]):
+        for segment in segments:
+            yield segment
 
-class GlinerFirstDateMethod(FirstDateMethod):
     @staticmethod
     def contains_year(text: str):
         year_pattern = re.compile(r"([0-9]{2})")
         return bool(year_pattern.search(text.replace(" ", "")))
 
-    def get_date_from_segments(self, segments: list[PdfDataSegment], languages):
+    def get_date_from_segments(self, model, segments: list[PdfDataSegment], languages):
         merge_segments: list[list[PdfDataSegment]] = self.merge_segments_for_dates(segments)
         for segments in merge_segments:
             segment_merged = PdfDataSegment.from_list_to_merge(segments)
             if not self.contains_year(segment_merged.text_content):
                 continue
 
-            date = GlinerDateParserMethod.get_date([segment_merged.text_content])
+            date = GlinerDateParserMethod.get_date(model, [segment_merged.text_content])
             if date:
                 for segment in segments:
                     segment.ml_label = 1

src/trainable_entity_extractor/adapters/extractors/pdf_to_text_extractor/methods/GlinerLastDateMethod.py

@@ -13,12 +13,12 @@ def loop_segments(segments):
         for segment in reversed(segments):
             yield segment
 
-    def get_date_from_segments(self, segments, languages):
+    def get_date_from_segments(self, model, segments, languages):
         for segment in self.loop_segments(segments):
             if not self.contains_year(segment.text_content):
                 continue
 
-            date = GlinerDateParserMethod.get_date([segment.text_content])
+            date = GlinerDateParserMethod.get_date(model, [segment.text_content])
             if date:
                 segment.ml_label = 1
                 return date.strftime("%Y-%m-%d")

src/trainable_entity_extractor/adapters/extractors/text_to_text_extractor/methods/GlinerDateParserMethod.py

@@ -1,3 +1,5 @@
+from gliner import GLiNER
+
 from trainable_entity_extractor.domain.ExtractionData import ExtractionData
 from trainable_entity_extractor.domain.PredictionSamplesData import PredictionSamplesData
 from trainable_entity_extractor.adapters.extractors.ToTextExtractorMethod import ToTextExtractorMethod
@@ -15,12 +17,12 @@ def get_alphanumeric_text_with_spaces(text):
         return "".join([letter for letter in text if letter.isalnum() or letter.isspace()])
 
     @staticmethod
-    def get_date(tags_texts: list[str]):
+    def get_date(model, tags_texts: list[str]):
         if not tags_texts:
             return ""
         text = GlinerDateParserMethod.get_alphanumeric_text_with_spaces(" ".join(tags_texts))
         try:
-            gliner_date_extractor = GlinerDateExtractor()
+            gliner_date_extractor = GlinerDateExtractor(model)
             dates = gliner_date_extractor.extract_dates(text)
             return dates[0]
         except:
@@ -29,7 +31,9 @@ def get_date(tags_texts: list[str]):
         return None
 
     def train(self, extraction_data: ExtractionData):
-        gliner_date_extractor = GlinerDateExtractor()
+        gliner_model = GLiNER.from_pretrained("urchade/gliner_multi-v2.1")
+
+        gliner_date_extractor = GlinerDateExtractor(gliner_model)
 
         for sample in extraction_data.samples[:15]:
             if not sample.labeled_data.label_text.strip():
@@ -42,11 +46,13 @@ def train(self, extraction_data: ExtractionData):
         self.save_json(self.IS_VALID_EXECUTION_FILE_NAME, "true")
 
     def predict(self, prediction_samples_data: PredictionSamplesData) -> list[str]:
+        gliner_model = GLiNER.from_pretrained("urchade/gliner_multi-v2.1")
+
         if self.load_json(self.IS_VALID_EXECUTION_FILE_NAME) == "false":
             return [""] * len(prediction_samples_data.prediction_samples)
 
         predictions_dates = [
-            self.get_date(prediction_sample.get_input_text_by_lines())
+            self.get_date(gliner_model, prediction_sample.get_input_text_by_lines())
             for prediction_sample in prediction_samples_data.prediction_samples
         ]
         predictions = [date.strftime("%Y-%m-%d") if date else "" for date in predictions_dates]