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")

# Variável global para armazenar o conteúdo do PDF
pdf_content = ""

# 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 carregar o PDF e armazenar o conteúdo
def load_pdf(pdf_file):
    global pdf_content
    pdf_content = read_pdf(pdf_file.name)
    return "PDF carregado com sucesso!"

# Função para responder perguntas com base no conteúdo do PDF
def answer_question(question, max_length=200, temperature=0.7, top_k=50, top_p=0.95):
    global pdf_content
    if not pdf_content:
        return "Por favor, carregue um PDF primeiro."
    prompt = f"Conteúdo do PDF: {pdf_content}\nPergunta: {question}\nResposta em português:"
    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 para carregar PDF e fazer perguntas
pdf_loader = gr.Interface(
    fn=load_pdf,
    inputs=gr.File(label="Carregue um PDF"),
    outputs="text",
    title="Carregar PDF"
)

question_answerer = gr.Interface(
    fn=answer_question,
    inputs=gr.Textbox(lines=2, label="Pergunta"),
    outputs="text",
    title="Perguntas sobre o PDF"
)

# Combinar as interfaces em uma aplicação
iface = gr.TabbedInterface(
    [pdf_loader, question_answerer],
    ["Carregar PDF", "Fazer Perguntas"]
)

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