import gradio as gr
import pdfplumber
import re
from transformers import pipeline

# Model do analizy układu faktur
extractor = pipeline("document-question-answering", model="impira/layoutlm-document-qa")

def extract_seller_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)

    # Model zwraca tekst, który uznał za odpowiedź na pytanie
    seller_info = response[0]["answer"] if response else "Nie znaleziono"

    return {"Sprzedawca": seller_info}

# Interfejs użytkownika w Hugging Face Spaces
iface = gr.Interface(
    fn=extract_seller_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."
)

if __name__ == "__main__":
    iface.launch()