Update app.py

app.py

@@ -1,29 +1,37 @@
 import gradio as gr
-from magic_pdf import MagicPDF
+import pdfplumber
+from transformers import pipeline
+
+# Inicjalizacja modelu do ekstrakcji informacji
+extractor = pipeline("ner", model="dslim/bert-base-NER")
 
 def extract_info(pdf_file):
-    # Inicjalizacja MagicPDF
-    pdf_processor = MagicPDF()
+    with pdfplumber.open(pdf_file) as pdf:
+        text = ""
+        for page in pdf.pages:
+            text += page.extract_text() + "\n"
 
-    # Przetwarzanie pliku PDF
-    result = pdf_processor(pdf_file.name)
+    # Przetwarzanie tekstu modelem NLP
+    entities = extractor(text)
 
-    # Wyodrębnianie pożądanych informacji
-    extracted_data = {
-        "company_name": result.get("company_name", "Nie znaleziono"),
-        "invoice_number": result.get("invoice_number", "Nie znaleziono"),
-        # Dodaj inne pola według potrzeb
-    }
+    # Filtrowanie i formatowanie wyników
+    extracted_data = {}
+    for entity in entities:
+        label = entity['entity']
+        word = entity['word']
+        if label not in extracted_data:
+            extracted_data[label] = []
+        extracted_data[label].append(word)
 
     return extracted_data
 
-# Interfejs użytkownika
+# Interfejs użytkownika w Hugging Face Space
 iface = gr.Interface(
     fn=extract_info,
-    inputs=gr.inputs.File(label="Wybierz plik PDF"),
-    outputs=gr.outputs.JSON(label="Wyodrębnione informacje"),
+    inputs=gr.File(label="Wybierz plik PDF"),
+    outputs="json",
     title="Ekstrakcja informacji z faktur PDF",
-    description="Prześlij plik PDF z fakturą, aby wyodrębnić określone informacje."
+    description="Prześlij plik PDF z fakturą, a model rozpozna kluczowe informacje."
 )
 
 if __name__ == "__main__":