import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel
import torch

# Modelos
base_model = "microsoft/phi-2"
model_peft = "vsvasconcelos/Phi-2_PT_QA_2_v3"

# Carrega o tokenizer
tokenizer = AutoTokenizer.from_pretrained(base_model)

# Carrega o modelo base e o modelo treinado com PEFT
model = AutoModelForCausalLM.from_pretrained(base_model)
model = PeftModel.from_pretrained(model, model_peft)

# Adicione uma mensagem de confirmação
print("Modelo PEFT vsvasconcelos/Phi-2_PT_QA_2_v3 carregado com sucesso!")

# Define a função para gerar respostas
def generate_response(model, tokenizer, prompt, max_length=512, num_return_sequences=1):
    # Tokenizando o prompt de entrada
    inputs = tokenizer(prompt, return_tensors='pt')

    model.eval()

    # Gerando a resposta usando o modelo
    with torch.no_grad():
        outputs = model.generate(
            inputs['input_ids'],
            max_length=max_length,
            num_return_sequences=num_return_sequences,
            pad_token_id=tokenizer.eos_token_id,
            eos_token_id=tokenizer.eos_token_id,
            early_stopping=True
        )

    # Tokenizar a resposta gerada
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return response

# Define a função para responder às perguntas
def answer_question(context, question):
    if context:
        prompt = f"Context: {context}\nQuestion: {question}"
    else:
        prompt = f"Question: {question}"
    return generate_response(model, tokenizer, prompt)

# Cria a interface Gradio
iface = gr.Interface(
    fn=answer_question,
    inputs=["text", "text"],
    outputs="text",
    title="QA Model",
    description="Pergunte algo e obtenha uma resposta!"
)

iface.launch()