Update app.py

app.py

@@ -3,48 +3,30 @@ import pdfplumber
 import re
 from transformers import pipeline
 
-# Model NER do rozpoznawania nazw organizacji
-extractor = pipeline("ner", model="dbmdz/bert-large-cased-finetuned-conll03-english", aggregation_strategy="simple")
+# Model do analizy układu faktur
+extractor = pipeline("document-question-answering", model="impira/layoutlm-document-qa")
 
-def extract_seller(pdf_file):
+def extract_seller_data(pdf_file):
     with pdfplumber.open(pdf_file) as pdf:
-        # Pobranie tekstu z PDF
+        # Pobranie całego 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, aby model działał szybciej)
-    chunks = [full_text[i:i+512] for i in range(0, len(full_text), 512)]
+    # Zadawanie pytania modelowi → pytamy o cały blok danych sprzedawcy
+    question = "What is the seller's information?"
+    response = extractor(question=question, context=full_text)
 
-    seller_tokens = []
+    # Model zwraca tekst, który uznał za odpowiedź na pytanie
+    seller_info = response[0]["answer"] if response else "Nie znaleziono"
 
-    for chunk in chunks:
-        entities = extractor(chunk)
-
-        for entity in entities:
-            if "ORG" in entity["entity_group"]:  # Szukamy nazw organizacji
-                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
-
-    # Łą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"}
+    return {"Sprzedawca": seller_info}
 
 # Interfejs użytkownika w Hugging Face Spaces
 iface = gr.Interface(
-    fn=extract_seller,
+    fn=extract_seller_data,
     inputs=gr.File(label="Wybierz plik PDF"),
     outputs="json",
-    title="Ekstrakcja Sprzedawcy z Faktury",
-    description="Prześlij plik PDF, aby wydobyć nazwę sprzedawcy."
+    title="Ekstrakcja danych sprzedawcy z faktury",
+    description="Prześlij plik PDF, a model zwróci kompletny zestaw danych o sprzedawcy."
 )
 
 if __name__ == "__main__":