Update app.py

app.py

@@ -1,35 +1,27 @@
 import gradio as gr
-import pdfplumber
-from transformers import pipeline
-
-# Inicjalizacja modelu
-extractor = pipeline("ner", model="opendatalab/PDF-Extract-Kit")
+from magic_pdf import MagicPDF
 
 def extract_info(pdf_file):
-    with pdfplumber.open(pdf_file) as pdf:
-        text = ""
-        for page in pdf.pages:
-            text += page.extract_text()
+    # Inicjalizacja MagicPDF
+    pdf_processor = MagicPDF()
 
-    # Przetwarzanie tekstu za pomocą modelu
-    entities = extractor(text)
+    # Przetwarzanie pliku PDF
+    result = pdf_processor(pdf_file.name)
 
-    # Filtrowanie i formatowanie wyników
-    results = {}
-    for entity in entities:
-        label = entity['entity']
-        word = entity['word']
-        if label not in results:
-            results[label] = []
-        results[label].append(word)
+    # 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
+    }
 
-    return results
+    return extracted_data
 
 # Interfejs użytkownika
 iface = gr.Interface(
     fn=extract_info,
     inputs=gr.inputs.File(label="Wybierz plik PDF"),
-    outputs=gr.outputs.JSON(label="Wykryte informacje"),
+    outputs=gr.outputs.JSON(label="Wyodrębnione informacje"),
     title="Ekstrakcja informacji z faktur PDF",
     description="Prześlij plik PDF z fakturą, aby wyodrębnić określone informacje."
 )