Update app.py

app.py

@@ -3,30 +3,50 @@ import pdfplumber
 import re
 from transformers import pipeline
 
-# Model do analizy układu faktur
-extractor = pipeline("document-question-answering", model="impira/layoutlm-document-qa")
+# Model do rozpoznawania nazw organizacji i wartości numerycznych
+extractor = pipeline("ner", model="dbmdz/bert-large-cased-finetuned-conll03-english", aggregation_strategy="simple")
 
-def extract_seller_data(pdf_file):
+# Reguły do identyfikacji wartości liczbowych (NIP, kwoty, daty)
+nip_pattern = re.compile(r'\b\d{10}\b')  # Polski NIP: 10 cyfr
+kwota_pattern = re.compile(r'\b\d+[\.,]?\d*\b')  # Kwoty: np. 123.45 lub 123
+data_pattern = re.compile(r'\b\d{2}\.\d{2}\.\d{4}\b')  # Daty: np. 21.10.2024
+
+def extract_invoice_data(pdf_file):
     with pdfplumber.open(pdf_file) as pdf:
-        # Pobranie całego tekstu z PDF
         full_text = "\n".join(page.extract_text() for page in pdf.pages if page.extract_text())
 
-    # Zadawanie pytania modelowi → pytamy o cały blok danych sprzedawcy
-    question = "What is the seller's information?"
-    response = extractor(question=question, context=full_text)
+    # Szukamy danych w tekście
+    entities = extractor(full_text)
+    
+    seller_name = []
+    seller_nip = None
+    items = []
+    total_amount = None
+    invoice_date = None
+
+    for entity in entities:
+        if "ORG" in entity["entity_group"]:
+            seller_name.append(entity["word"])  # Zbieramy nazwę sprzedawcy
 
-    # Model zwraca tekst, który uznał za odpowiedź na pytanie
-    seller_info = response[0]["answer"] if response else "Nie znaleziono"
+    # Znajdujemy wartości numeryczne dla NIP, kwot, dat
+    seller_nip = nip_pattern.search(full_text)
+    total_amount = max(kwota_pattern.findall(full_text), key=float, default=None)
+    invoice_date = data_pattern.search(full_text)
 
-    return {"Sprzedawca": seller_info}
+    return {
+        "Sprzedawca": " ".join(seller_name) if seller_name else "Nie znaleziono",
+        "NIP": seller_nip.group() if seller_nip else "Nie znaleziono",
+        "Data faktury": invoice_date.group() if invoice_date else "Nie znaleziono",
+        "Kwota całkowita": total_amount if total_amount else "Nie znaleziono"
+    }
 
 # Interfejs użytkownika w Hugging Face Spaces
 iface = gr.Interface(
-    fn=extract_seller_data,
+    fn=extract_invoice_data,
     inputs=gr.File(label="Wybierz plik PDF"),
     outputs="json",
-    title="Ekstrakcja danych sprzedawcy z faktury",
-    description="Prześlij plik PDF, a model zwróci kompletny zestaw danych o sprzedawcy."
+    title="Ekstrakcja danych z faktury",
+    description="Prześlij plik PDF, a model zwróci dane sprzedawcy, NIP, kwotę i datę faktury."
 )
 
 if __name__ == "__main__":