Add data summary

Author: gabriel-piles

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

@@ -57,7 +57,7 @@ def get_performance(self, train_set: ExtractionData, test_set: ExtractionData) -
 
         self.train(train_set)
         prediction_samples_data = PredictionSamplesData(
-            prediction_samples=[PredictionSample.from_text(x.pdf_data.get_text()) for x in test_set.samples],
+            prediction_samples=[PredictionSample.from_pdf_data(x.pdf_data) for x in test_set.samples],
             options=self.options,
             multi_value=self.multi_value,
         )

src/trainable_entity_extractor/adapters/extractors/segment_selector/methods/lightgbm_frequent_words/LightgbmFrequentWords.py

@@ -30,7 +30,7 @@ def __init__(self):
 
     def create_model(self, training_pdfs_segments: list[PdfData], model_path):
         start = time()
-        self.set_segments(pdfs_segments=training_pdfs_segments)
+        self.set_segments(pdfs_data=training_pdfs_segments)
 
         config_logger.info(f"Set segments {int(time() - start)} seconds")
 
@@ -81,13 +81,13 @@ def get_training_data(self):
 
         return X, y
 
-    def set_segments(self, pdfs_segments: list[PdfData]):
+    def set_segments(self, pdfs_data: list[PdfData]):
         self.segments = list()
-        for pdf_features in pdfs_segments:
-            self.segments.extend(SegmentLightgbmFrequentWords.from_pdf_features(pdf_features))
+        for pdf_data in pdfs_data:
+            self.segments.extend(SegmentLightgbmFrequentWords.from_pdf_data(pdf_data))
 
-    def predict(self, model, testing_pdfs_segments: list[PdfData], model_path):
-        self.set_segments(testing_pdfs_segments)
+    def predict(self, model, pdfs_data: list[PdfData], model_path):
+        self.set_segments(pdfs_data)
         self.set_most_frequent_words_to_segments(model_path)
         x, y = self.get_training_data()
         x = x[:, : model.num_feature()]

src/trainable_entity_extractor/adapters/extractors/segment_selector/methods/lightgbm_frequent_words/SegmentLightgbmFrequentWords.py

@@ -235,11 +235,11 @@ def token_after_last_token(self, token: PdfToken):
         return False
 
     @staticmethod
-    def from_pdf_features(pdf_features: PdfData) -> list["SegmentLightgbmFrequentWords"]:
-        modes = Modes(pdf_features)
+    def from_pdf_data(pdf_data: PdfData) -> list["SegmentLightgbmFrequentWords"]:
+        modes = Modes(pdf_data)
         segments: list["SegmentLightgbmFrequentWords"] = list()
-        for index, pdf_segment in enumerate(pdf_features.pdf_data_segments):
-            segment_landmarks = SegmentLightgbmFrequentWords(index, pdf_segment, pdf_features, modes)
+        for index, pdf_segment in enumerate(pdf_data.pdf_data_segments):
+            segment_landmarks = SegmentLightgbmFrequentWords(index, pdf_segment, pdf_data, modes)
             segments.append(segment_landmarks)
 
         sorted_pdf_segments = sorted(segments, key=lambda x: (x.page_index, x.top))

src/trainable_entity_extractor/domain/ExtractionDataSummary.py

@@ -0,0 +1,174 @@
+from collections import Counter
+from typing import Optional
+
+from pydantic import BaseModel, Field
+
+from trainable_entity_extractor.domain.ExtractionData import ExtractionData
+
+
+class OptionDistribution(BaseModel):
+    option_id: str
+    option_label: str
+    count: int
+    percentage: float
+
+
+class LanguageDistribution(BaseModel):
+    language_iso: str
+    count: int
+    percentage: float
+
+
+class TextLengthStats(BaseModel):
+    min_length: int
+    max_length: int
+    avg_length: float
+    median_length: float
+
+
+class ExtractionDataSummary(BaseModel):
+    total_samples: int
+    total_options: int
+    has_pdf_data: bool
+    empty_pdfs_count: int = 0
+    languages: list[LanguageDistribution] = Field(default_factory=list)
+    option_distribution: list[OptionDistribution] = Field(default_factory=list)
+    label_text_stats: Optional[TextLengthStats] = None
+    source_text_stats: Optional[TextLengthStats] = None
+    samples_with_values: int = 0
+
+    @staticmethod
+    def from_extraction_data(extraction_data: ExtractionData) -> "ExtractionDataSummary":
+        total_samples = len(extraction_data.samples)
+        total_options = len(extraction_data.options) if extraction_data.options else 0
+
+        has_pdf_data = any(sample.pdf_data and sample.pdf_data.get_text() for sample in extraction_data.samples)
+        empty_pdfs_count = 0
+
+        if has_pdf_data:
+            for sample in extraction_data.samples:
+                if sample.pdf_data:
+                    if not sample.pdf_data.get_text():
+                        empty_pdfs_count += 1
+
+        language_counter = Counter()
+        for sample in extraction_data.samples:
+            if sample.labeled_data and sample.labeled_data.language_iso:
+                language_counter[sample.labeled_data.language_iso] += 1
+
+        languages = [
+            LanguageDistribution(language_iso=lang, count=count, percentage=round(count / total_samples * 100, 2))
+            for lang, count in language_counter.most_common()
+        ]
+
+        option_counter = Counter()
+        for sample in extraction_data.samples:
+            if sample.labeled_data and sample.labeled_data.values:
+                for value in sample.labeled_data.values:
+                    option_counter[value.id] += 1
+
+        option_distribution = []
+        if extraction_data.options:
+            for option in extraction_data.options:
+                count = option_counter.get(option.id, 0)
+                option_distribution.append(
+                    OptionDistribution(
+                        option_id=option.id,
+                        option_label=option.label,
+                        count=count,
+                        percentage=round(count / total_samples * 100, 2) if total_samples > 0 else 0,
+                    )
+                )
+            option_distribution = sorted(option_distribution, key=lambda x: x.count, reverse=True)[:30]
+
+        label_text_lengths = []
+        source_text_lengths = []
+        samples_with_values = 0
+
+        for sample in extraction_data.samples:
+            if sample.labeled_data:
+                if sample.labeled_data.label_text:
+                    label_text_lengths.append(len(sample.labeled_data.label_text))
+                if sample.labeled_data.source_text:
+                    source_text_lengths.append(len(sample.labeled_data.source_text))
+                if sample.labeled_data.values:
+                    samples_with_values += 1
+
+        label_text_stats = None
+        if label_text_lengths:
+            sorted_lengths = sorted(label_text_lengths)
+            label_text_stats = TextLengthStats(
+                min_length=min(label_text_lengths),
+                max_length=max(label_text_lengths),
+                avg_length=round(sum(label_text_lengths) / len(label_text_lengths), 2),
+                median_length=sorted_lengths[len(sorted_lengths) // 2],
+            )
+
+        source_text_stats = None
+        if source_text_lengths:
+            sorted_lengths = sorted(source_text_lengths)
+            source_text_stats = TextLengthStats(
+                min_length=min(source_text_lengths),
+                max_length=max(source_text_lengths),
+                avg_length=round(sum(source_text_lengths) / len(source_text_lengths), 2),
+                median_length=sorted_lengths[len(sorted_lengths) // 2],
+            )
+
+        return ExtractionDataSummary(
+            total_samples=total_samples,
+            total_options=total_options,
+            has_pdf_data=has_pdf_data,
+            empty_pdfs_count=empty_pdfs_count,
+            languages=languages,
+            option_distribution=option_distribution,
+            label_text_stats=label_text_stats,
+            source_text_stats=source_text_stats,
+            samples_with_values=samples_with_values,
+        )
+
+    def to_report_string(self) -> str:
+        lines = [
+            "Data Summary",
+            "=" * 80,
+            f"Total Samples: {self.total_samples}",
+        ]
+
+        if self.total_options:
+            lines.append(f"Total Options: {self.total_options}")
+
+        if self.total_options and self.option_distribution:
+            lines.append("\nOption Distribution:")
+            for dist in self.option_distribution:
+                lines.append(f"  - {dist.option_label} (id: {dist.option_id}): {dist.count} samples ({dist.percentage}%)")
+
+        if self.samples_with_values > 0:
+            percentage = round(self.samples_with_values / self.total_samples * 100, 2)
+            lines.append(f"\nSamples with Option Values: {self.samples_with_values} ({percentage}%)")
+
+        if self.languages:
+            lines.append("\nLanguage Distribution:")
+            for lang_dist in self.languages:
+                lines.append(f"  - {lang_dist.language_iso}: {lang_dist.count} samples ({lang_dist.percentage}%)")
+
+        if self.has_pdf_data:
+            lines.append(f"\nPDF Data: Present")
+            if self.empty_pdfs_count > 0:
+                lines.append(f"Empty PDFs: {self.empty_pdfs_count}")
+
+        if self.label_text_stats:
+            lines.append("\nLabel Text Length:")
+            lines.append(f"  - Min: {self.label_text_stats.min_length}")
+            lines.append(f"  - Max: {self.label_text_stats.max_length}")
+            lines.append(f"  - Average: {self.label_text_stats.avg_length}")
+            lines.append(f"  - Median: {self.label_text_stats.median_length}")
+
+        if self.source_text_stats:
+            lines.append("\nSource Text Length:")
+            lines.append(f"  - Min: {self.source_text_stats.min_length}")
+            lines.append(f"  - Max: {self.source_text_stats.max_length}")
+            lines.append(f"  - Average: {self.source_text_stats.avg_length}")
+            lines.append(f"  - Median: {self.source_text_stats.median_length}")
+
+        lines.append("=" * 80)
+
+        return "\n".join(lines)

src/trainable_entity_extractor/domain/PerformanceSummary.py

@@ -17,7 +17,6 @@ class PerformanceSummary(BaseModel):
     performances: list[PerformanceLog] = []
     extraction_identifier: ExtractionIdentifier | None = None
     previous_timestamp: int = Field(default_factory=lambda: int(time()))
-    empty_pdf_count: int = 0
 
     def add_performance(self, method_name: str, performance: float, failed: bool = False):
         current_time = int(time())
@@ -53,7 +52,6 @@ def to_log(self) -> str:
         text += f"Samples: {self.samples_count}\n"
         text += f"Train/test: {self.training_samples_count}/{self.testing_samples_count}\n"
         text += f"{len(self.languages)} language(s): {', '.join(self.languages) if self.languages else 'None'}\n"
-        text += f"Empty PDFs: {self.empty_pdf_count}\n" if self.empty_pdf_count else ""
         text += f"Options count: {self.options_count}\n" if self.options_count > 0 else ""
         text += "Methods by performance:\n"
         for performance in sorted(self.performances, key=lambda x: x.performance, reverse=True):
@@ -94,5 +92,4 @@ def from_distributed_job(distributed_job: DistributedJob) -> "PerformanceSummary
             languages=languages,
             training_samples_count=training_samples_count,
             testing_samples_count=testing_samples_count,
-            empty_pdf_count=0,
         )

src/trainable_entity_extractor/domain/PredictionSample.py

@@ -32,8 +32,10 @@ def get_input_text_by_lines(self) -> list[str]:
         return [""]
 
     @staticmethod
-    def from_pdf_data(pdf_data: PdfData):
-        return PredictionSample(pdf_data=pdf_data)
+    def from_pdf_data(pdf_data: PdfData, entity_name: str = ""):
+        prediction_sample = PredictionSample.from_text(pdf_data.get_text(), entity_name)
+        prediction_sample.pdf_data = pdf_data
+        return prediction_sample
 
     @staticmethod
     def from_text(text: str, entity_name: str = ""):

src/trainable_entity_extractor/domain/Suggestion.py

@@ -108,8 +108,18 @@ def from_prediction_text(extraction_identifier: ExtractionIdentifier, entity_nam
     def from_prediction_multi_option(extraction_identifier: ExtractionIdentifier, entity_name: str, values: list[Value]):
         suggestion = Suggestion.get_empty(extraction_identifier, entity_name)
         suggestion.values = values
-        if values:
-            suggestion.segment_text = values[0].segment_text
+        for value in values:
+            if value.segment_text:
+                suggestion._raw_context = [values[0].segment_text]
+                suggestion.segment_text = FormatSegmentText([values[0].segment_text], value.label).get_text()
+                break
+
+        for value in values:
+            if value.segment_text:
+                value.segment_text = FormatSegmentText([value.segment_text], value.label).get_text()
+            else:
+                value.segment_text = FormatSegmentText(suggestion._raw_context, value.label).get_text()
+
         return suggestion
 
     def set_segment_text_from_sample(self, prediction_sample: PredictionSample):

src/trainable_entity_extractor/domain/Value.py

@@ -17,7 +17,7 @@ def __eq__(self, other):
         if not isinstance(other, Value):
             return False
 
-        if other.segment_text and self.segment_text != other.segment_text:
+        if self.segment_text and other.segment_text and self.segment_text != other.segment_text:
             return False
 
         return self.id == other.id and self.label == other.label

src/trainable_entity_extractor/tests/unit_tests/domain/test_PerformanceSummary.py

@@ -125,14 +125,12 @@ def test_direct_instantiation_with_empty_pdf_count(self):
             languages=["en", "es"],
             training_samples_count=25,
             testing_samples_count=15,
-            empty_pdf_count=2,
         )
 
         assert result.extractor_name == "PDF Extractor"
         assert result.samples_count == 4
         assert result.options_count == 0
         assert set(result.languages) == {"en", "es"}
-        assert result.empty_pdf_count == 2
 
     def test_to_log_basic_summary_no_methods(self):
         """Test to_log with basic summary but no performance methods"""

src/trainable_entity_extractor/tests/unit_tests/extractors/pdf_to_multi_option_extractor/multi_labels_methods/test_single_label_setfit_english.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.SingleLabelSetFitEnglishMethod import (
     SingleLabelSetFitEnglishMethod,
@@ -25,7 +27,7 @@ class TestSetFitSingleLabelEnglishMethod(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
@@ -45,21 +47,19 @@ def test_train_and_predict(self):
         extraction_data = ExtractionData(
             multi_value=False, options=options, samples=samples, extraction_identifier=extraction_identifier
         )
-        setfit_english_method = SingleLabelSetFitEnglishMethod(extraction_identifier, options, False)
+        setfit_english_method = SingleLabelSetFitEnglishMethod(extraction_identifier)
 
         try:
             setfit_english_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_3 = TrainingSample(pdf_data=pdf_data_3)
+        prediction_sample_1 = PredictionSample(pdf_data=pdf_data_1)
+        prediction_sample_2 = PredictionSample(pdf_data=pdf_data_2)
+        prediction_sample_3 = PredictionSample(pdf_data=pdf_data_3)
         prediction_samples = [prediction_sample_1, prediction_sample_2, prediction_sample_3]
 
-        prediction_data = ExtractionData(
-            multi_value=False, options=options, samples=prediction_samples, extraction_identifier=extraction_identifier
-        )
+        prediction_data = PredictionSamplesData(multi_value=False, options=options, prediction_samples=prediction_samples)
         predictions = setfit_english_method.predict(prediction_data)
 
         self.assertEqual(3, len(predictions))