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)

#Uma maquiagem pra deixar a resposta com a primeira letra em maiusculo
def capitalize(text):
    if text:
        return text[0].upper() + text[1:]
    return text



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

    # Extrair a parte relevante da resposta
    if "### Resposta:" in response:
        response = response.split("### Resposta:")[-1].strip()

    # Capitalizar a primeira letra da resposta
    response = capitalize(response)
    
    return response
    

# Define a função para responder às perguntas
def answer_question(context, question):
    prompt = f"""Com base no contexto fornecido, responda a questão de forma sucinta.
    Contexto:
    {context}
    Questão:
    {question}
    ### Resposta:
    """
    #else:
        #prompt = f"Questão:\n{question}\n### Resposta:"
    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()