import gradio as gr
from transformers import pipeline
import PyPDF2

# Carregar o modelo de linguagem gratuito da Hugging Face
nlp = pipeline("question-answering", model="distilbert-base-cased-distilled-squad")

def extract_text_from_pdf(pdf_file):
    with open(pdf_file.name, "rb") as file:
        reader = PyPDF2.PdfFileReader(file)
        text = ""
        for page_num in range(reader.numPages):
            page = reader.getPage(page_num)
            text += page.extract_text()
    return text

def answer_question(pdf_file, question):
    # Extrair texto do PDF
    context = extract_text_from_pdf(pdf_file)

    # Usar o modelo para responder a pergunta
    result = nlp(question=question, context=context)
    return result['answer']

# Interface Gradio
iface = gr.Interface(
    fn=answer_question,
    inputs=[
        gr.inputs.File(label="Carregar PDF"),
        gr.inputs.Textbox(label="Pergunta")
    ],
    outputs=gr.outputs.Textbox(label="Resposta"),
    title="QA sobre PDF",
    description="Carregue um PDF e faça perguntas sobre o conteúdo."
)

# Iniciar a interface
iface.launch()