import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
import fitz  # PyMuPDF

# Carregar o tokenizer e o modelo
tokenizer = AutoTokenizer.from_pretrained("Locutusque/gpt2-xl-conversational")
model = AutoModelForCausalLM.from_pretrained("Locutusque/gpt2-xl-conversational")

# Função para ler o PDF
def read_pdf(file_path):
    doc = fitz.open(file_path)
    text = ""
    for page_num in range(len(doc)):
        page = doc.load_page(page_num)
        text += page.get_text()
    return text

# Função para gerar respostas a partir do texto do PDF
def generate_response_from_pdf(pdf_file, max_length=200, temperature=0.7, top_k=50, top_p=0.95):
    # Lê o conteúdo do PDF
    pdf_text = read_pdf(pdf_file.name)
    # Gerar prompt em português
    prompt = f"Texto do PDF: {pdf_text}\nResposta em português:"
    # Gera a resposta usando o texto do PDF como prompt
    inputs = tokenizer(prompt, return_tensors="pt")
    outputs = model.generate(
        inputs.input_ids,
        max_length=max_length,
        temperature=temperature,
        top_k=top_k,
        top_p=top_p,
        num_return_sequences=1
    )
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return response

# Interface do Gradio
iface = gr.Interface(
    fn=generate_response_from_pdf,
    inputs=gr.inputs.File(label="Carregue um PDF"),
    outputs="text",
    title="GPT-2 Conversational com PDF em Português"
)

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