Update app.py

app.py

@@ -2,11 +2,12 @@ import gradio as gr
 import pdfplumber
 import re
 from transformers import LayoutLMForTokenClassification, AutoTokenizer
+import torch
 
 # Wczytanie modelu LayoutLMv3
 model_name = "kryman27/layoutlmv3-finetuned"
 model = LayoutLMForTokenClassification.from_pretrained(model_name)
-tokenizer = AutoTokenizer.from_pretrained(model_name)  # Poprawiona wersja
+tokenizer = AutoTokenizer.from_pretrained(model_name)  # Automatyczne wykrycie tokenizatora
 
 # Reguły do wykrywania NIP, kwot, dat
 nip_pattern = re.compile(r'\bPL\s?\d{10}\b|\b\d{10}\b')
@@ -16,37 +17,45 @@ payment_keywords = ["data płatności", "termin płatności", "zapłata", "płat
 
 def extract_invoice_data(pdf_file):
     with pdfplumber.open(pdf_file) as pdf:
-        full_text = "\n".join(page.extract_text() for page in pdf.pages if page.extract_text())
+        words, boxes = [], []
 
-    # Podział tekstu na listę słów (LayoutLMv3 wymaga tokenizacji na poziomie słów)
-    words = full_text.split()  # Nowa poprawiona linia
-    tokens = tokenizer(words, is_split_into_words=True, return_tensors="pt", truncation=True)  # Poprawiona linia
+        for page in pdf.pages:
+            extracted_words = page.extract_words()
+            for word in extracted_words:
+                words.append(word['text'])  # Pobieramy tekst słowa
+                bbox = [word['x0'], word['top'], word['x1'], word['bottom']]
+                boxes.append(bbox)  # Pobieramy bounding box (pozycję słowa na stronie)
+
+    # Tokenizacja tekstu + dodanie bounding boxes
+    tokens = tokenizer(words, boxes=boxes, is_split_into_words=True, return_tensors="pt", truncation=True)
 
     # Predykcja modelu
-    outputs = model(**tokens)
+    with torch.no_grad():
+        outputs = model(**tokens)
     predictions = outputs.logits.argmax(-1).squeeze().tolist()
 
     # Przetwarzanie wyników
     entities = []
-    for token, pred in zip(words, predictions):  # Teraz iterujemy po `words`
+    for token, pred in zip(words, predictions):
         if pred > 0:  # Pomijamy tło
             entities.append((token, model.config.id2label[pred]))
 
     # Wyszukiwanie kluczowych wartości
     seller_name = [token for token, label in entities if "ORG" in label]
-    seller_nip = nip_pattern.search(full_text)
-    kwoty = kwota_pattern.findall(full_text)
+    seller_nip = nip_pattern.search(" ".join(words))
+    kwoty = kwota_pattern.findall(" ".join(words))
     kwoty = [float(k.replace(",", ".")) for k in kwoty if k.replace(",", ".").replace(".", "").isdigit()]
     total_amount = max(kwoty) if kwoty else None
 
     # Szukamy daty płatności
     payment_date = None
-    for line in full_text.split("\n"):
-        if any(keyword in line.lower() for keyword in payment_keywords):
-            date_match = data_pattern.search(line)
-            if date_match:
-                payment_date = date_match.group()
-                break
+    for i, word in enumerate(words):
+        if any(keyword in word.lower() for keyword in payment_keywords):
+            if i + 1 < len(words):
+                date_match = data_pattern.search(words[i + 1])
+                if date_match:
+                    payment_date = date_match.group()
+                    break
 
     return {
         "Sprzedawca": " ".join(seller_name) if seller_name else "Nie znaleziono",