import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer

# Carregar o modelo e o tokenizer
modelo = AutoModelForCausalLM.from_pretrained("unicamp-dl/ptt5-base-portuguese-vocab")
tokenizer = AutoTokenizer.from_pretrained("unicamp-dl/ptt5-base-portuguese-vocab")

# Função para gerar resposta
def gerar_resposta(texto_entrada):
    inputs = tokenizer(texto_entrada, return_tensors="pt", padding=True)
    outputs = modelo.generate(**inputs, max_length=100)
    resposta = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return resposta

# Interface Gradio
iface = gr.Interface(
    fn=gerar_resposta,
    inputs=gr.Textbox(lines=2, placeholder="Digite sua mensagem aqui..."),
    outputs="text",
    title="Chatbot em Português",
    description="Um chatbot simples usando PTT5."
)

iface.launch()