import gradio as gr
from transformers import pipeline
from sentence_transformers import SentenceTransformer, util

generator_a = pipeline("text-generation", model="gpt2")
generator_b = pipeline("text2text-generation", model="google/flan-t5-base")
similarity_model = SentenceTransformer("sentence-transformers/paraphrase-MiniLM-L6-v2")

def comparar(prompt):
    resp_a = generator_a(prompt, max_new_tokens=60, temperature=0.7)[0]["generated_text"]
    resp_b = generator_b(prompt, max_new_tokens=60, temperature=0.7)[0]["generated_text"]

    emb_a = similarity_model.encode(resp_a, convert_to_tensor=True)
    emb_b = similarity_model.encode(resp_b, convert_to_tensor=True)
    similaridade = util.cos_sim(emb_a, emb_b).item()

    return resp_a.strip(), resp_b.strip(), f"{similaridade:.4f}"

gr.Interface(
    fn=comparar,
    inputs=gr.Textbox(label="Digite um prompt"),
    outputs=[
        gr.Textbox(label="Resposta do GPT-2"),
        gr.Textbox(label="Resposta do Flan-T5"),
        gr.Textbox(label="Similaridade entre respostas")
    ],
    title="Comparador de Modelos LLM Leves"
).launch()