import gradio as gr
import pdfplumber
from transformers import pipeline

# Inicjalizacja modelu LayoutLM do ekstrakcji tekstu
extractor = pipeline("ner", model="microsoft/layoutlmv3-base", aggregation_strategy="simple")

def extract_seller(pdf_file):
    with pdfplumber.open(pdf_file) as pdf:
        text = "\n".join(page.extract_text() for page in pdf.pages if page.extract_text())

    # Przetwarzanie tekstu modelem NLP
    entities = extractor(text)

    seller_name = None

    for entity in entities:
        if "ORG" in entity["entity_group"]:  # Szukamy nazw organizacji
            seller_name = entity["word"]
            break  # Pobieramy pierwszą wykrytą firmę jako sprzedawcę

    return {"Sprzedawca": seller_name if seller_name else "Nie znaleziono"}

# Interfejs użytkownika w Hugging Face Spaces
iface = gr.Interface(
    fn=extract_seller,
    inputs=gr.File(label="Wybierz plik PDF"),
    outputs="json",
    title="Ekstrakcja Sprzedawcy z Faktury",
    description="Prześlij plik PDF, aby wydobyć nazwę sprzedawcy."
)

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