Update app.py

app.py

@@ -1,5 +1,6 @@
 import gradio as gr
 import pdfplumber
+import re
 from transformers import pipeline
 
 # Model NER do rozpoznawania nazw organizacji
@@ -10,7 +11,7 @@ def extract_seller(pdf_file):
         # Pobranie tekstu z PDF
         full_text = "\n".join(page.extract_text() for page in pdf.pages if page.extract_text())
 
-    # Podział tekstu na krótkie fragmenty (maks. 512 znaków)
+    # Podział tekstu na krótkie fragmenty (maks. 512 znaków, aby model działał szybciej)
     chunks = [full_text[i:i+512] for i in range(0, len(full_text), 512)]
 
     seller_tokens = []
@@ -20,7 +21,11 @@ def extract_seller(pdf_file):
 
         for entity in entities:
             if "ORG" in entity["entity_group"]:  # Szukamy nazw organizacji
-                seller_tokens.append(entity["word"])
+                word = entity["word"]
+
+                # Usuwamy błędne tokeny (np. "Faktura", "Z", itp.)
+                if not re.match(r"^(Faktura|Z|##|I|KRZYSZTOF|ŻARNOWIECKA)$", word, re.IGNORECASE):
+                    seller_tokens.append(word)
 
         if seller_tokens:  # Jeśli znaleziono organizację, przerywamy pętlę
             break
@@ -28,6 +33,9 @@ def extract_seller(pdf_file):
     # Łączymy tokeny w pełną nazwę organizacji
     seller_name = " ".join(seller_tokens)
 
+    # Usuwamy ewentualne powtórzone litery lub spacje między literami np. "S A" → "S.A."
+    seller_name = re.sub(r"\bS A\b", "S.A.", seller_name)
+
     return {"Sprzedawca": seller_name if seller_name else "Nie znaleziono"}
 
 # Interfejs użytkownika w Hugging Face Spaces