# app.py

from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
from sentence_transformers import SentenceTransformer, util
import torch
import gradio as gr

# Modelos A e B
modelo_a_nome = "pierreguillou/t5-base-pt-small-finetuned-question-generation"
modelo_b_nome = "unicamp-dl/ptt5-base-portuguese-vocab"

tokenizer_a = AutoTokenizer.from_pretrained(modelo_a_nome)
modelo_a = AutoModelForSeq2SeqLM.from_pretrained(modelo_a_nome)

tokenizer_b = AutoTokenizer.from_pretrained(modelo_b_nome)
modelo_b = AutoModelForSeq2SeqLM.from_pretrained(modelo_b_nome)

# Modelo Árbitro
modelo_arbitro_nome = "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"
modelo_arbitro = SentenceTransformer(modelo_arbitro_nome)

# Dispositivo
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
modelo_a = modelo_a.to(device)
modelo_b = modelo_b.to(device)
modelo_arbitro = modelo_arbitro.to(device)

# Funções
def gerar_resposta_modelo_a(pergunta):
    entrada = tokenizer_a.encode(pergunta, return_tensors="pt").to(device)
    saida_ids = modelo_a.generate(entrada, max_length=50, num_beams=4, early_stopping=True)
    resposta = tokenizer_a.decode(saida_ids[0], skip_special_tokens=True)
    return resposta

def gerar_resposta_modelo_b(pergunta):
    entrada = tokenizer_b.encode(pergunta, return_tensors="pt").to(device)
    saida_ids = modelo_b.generate(entrada, max_length=50, num_beams=4, early_stopping=True)
    resposta = tokenizer_b.decode(saida_ids[0], skip_special_tokens=True)
    return resposta

def escolher_melhor_resposta(pergunta, resposta_a, resposta_b):
    embeddings = modelo_arbitro.encode([pergunta, resposta_a, resposta_b], convert_to_tensor=True)
    similaridade_a = util.pytorch_cos_sim(embeddings[0], embeddings[1]).item()
    similaridade_b = util.pytorch_cos_sim(embeddings[0], embeddings[2]).item()

    if similaridade_a > similaridade_b:
        melhor = "A"
        resposta_escolhida = resposta_a
    else:
        melhor = "B"
        resposta_escolhida = resposta_b

    return resposta_escolhida, melhor, similaridade_a, similaridade_b

def responder(pergunta):
    resposta_a = gerar_resposta_modelo_a(pergunta)
    resposta_b = gerar_resposta_modelo_b(pergunta)
    resposta_final, escolhido, sim_a, sim_b = escolher_melhor_resposta(pergunta, resposta_a, resposta_b)

    return {
        "Melhor Resposta": resposta_final,
        "Resposta do Modelo A": f"{resposta_a} (similaridade: {sim_a:.2f})",
        "Resposta do Modelo B": f"{resposta_b} (similaridade: {sim_b:.2f})",
        "Modelo Escolhido": f"Modelo {escolhido}"
    }

interface = gr.Interface(
    fn=responder,
    inputs=gr.Textbox(label="Digite sua pergunta"),
    outputs=[
        gr.Textbox(label="Melhor Resposta"),
        gr.Textbox(label="Resposta do Modelo A"),
        gr.Textbox(label="Resposta do Modelo B"),
        gr.Textbox(label="Modelo Escolhido"),
    ],
    title="Chatbot em Cascata",
    description="Este chatbot compara duas respostas geradas por modelos diferentes e escolhe a melhor com base na similaridade semântica. Respostas 100% em português."
)

interface.launch()