app.py

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

MODEL_NAME = "microsoft/DialoGPT-small"

# Carrega tokenizer e modelo uma única vez
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)

device = torch.device("cpu")
model.to(device)

def gerar_resposta(mensagem, historico):
    """
    Função chamada pelo ChatInterface.
    - mensagem: texto digitado pelo usuário
    - historico: lista de tuplas (user, bot) que o Gradio mantém
    A função DEVE retornar apenas a resposta (string).
    """

    if historico is None:
        historico = []

    mensagem = mensagem.strip()
    if not mensagem:
        return "Me manda uma pergunta primeiro? :)"

    # Monta um contexto simples com base no histórico
    conversa = ""
    for user_msg, bot_msg in historico:
        conversa += f"Usuário: {user_msg}\nBot: {bot_msg}\n"
    conversa += f"Usuário: {mensagem}\nBot:"

    # Tokeniza e gera texto com o modelo
    inputs = tokenizer(
        conversa,
        return_tensors="pt",
        truncation=True,
        max_length=512
    ).to(device)

    with torch.no_grad():
        output_ids = model.generate(
            **inputs,
            max_new_tokens=80,
            do_sample=True,
            top_p=0.9,
            temperature=0.7,
            pad_token_id=tokenizer.eos_token_id,
        )

    texto_gerado = tokenizer.decode(output_ids[0], skip_special_tokens=True)

    # Pega apenas o trecho depois do último "Bot:"
    if "Bot:" in texto_gerado:
        resposta = texto_gerado.split("Bot:")[-1].strip()
    else:
        resposta = texto_gerado.strip()

    if not resposta:
        resposta = (
            "Estou pensando aqui, mas não consegui formular uma resposta melhor ainda. "
            "Tenta perguntar de outro jeito? :)"
        )

    # IMPORTANTE: retornamos APENAS a resposta (string)
    return resposta


descricao = """
### Chatbot com IA Real (Transformers) 🤖

Este Space demonstra um chatbot usando um modelo de linguagem da biblioteca **Transformers**:
- Modelo: `microsoft/DialoGPT-small`
- Backend: `transformers` + `torch`
- Interface: **Gradio ChatInterface**

Não precisa de chave de API externa e roda em CPU no Hugging Face Spaces.
"""

iface = gr.ChatInterface(
    fn=gerar_resposta,
    title="Chatbot com Transformers - Demonstração",
    description=descricao,
    examples=[
        "Oi, quem é você?",
        "O que é Inteligência Artificial?",
        "Me dá uma dica para estudar programação?",
        "Como posso aprender sobre segurança da informação?",
    ],
    cache_examples=False,   # evita execução automática na inicialização
)

if __name__ == "__main__":
    iface.launch()