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

# Modelo pequeno, viável em CPU grátis do Hugging Face
MODEL_NAME = "microsoft/DialoGPT-small"

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME).to("cpu")


def responder(pergunta: str) -> str:
    """
    Bot de Perguntas e Respostas sobre Deep Learning.
    Recebe uma pergunta em texto e devolve uma resposta em português,
    com tom de professor explicando de forma simples.
    """

    pergunta = (pergunta or "").strip()
    if not pergunta:
        return "Me manda uma pergunta sobre Deep Learning 🙂"

    # Prompt de sistema: define o "personagem"
    system_prefix = (
        "Você é o Professor DL, um professor de Deep Learning. "
        "Responda SEMPRE em português do Brasil, de forma simples, didática e objetiva, "
        "usando exemplos práticos quando possível. "
        "Explique conceitos como redes neurais, camadas, CNN, RNN, overfitting, "
        "regularização, dropout, etc., sem fórmulas muito pesadas."
    )

    prompt = (
        system_prefix
        + "\n\n"
        + f"Aluno: {pergunta}\n"
        + "Professor DL:"
    )

    # Tokenização
    inputs = tokenizer(
        prompt,
        return_tensors="pt",
        truncation=True,
        max_length=512,
    )

    # Geração
    with torch.no_grad():
        output_ids = model.generate(
            **inputs,
            max_new_tokens=160,
            do_sample=True,
            top_p=0.9,
            temperature=0.7,
            pad_token_id=tokenizer.eos_token_id,
        )

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

    # Extrai só a parte depois de "Professor DL:"
    if "Professor DL:" in saida:
        resposta = saida.split("Professor DL:")[-1].strip()
    else:
        resposta = saida.strip()

    if not resposta:
        resposta = (
            "Boa pergunta! Tenta reformular ou ser um pouco mais específico "
            "sobre o que você quer saber em Deep Learning."
        )

    return resposta


demo = gr.Interface(
    fn=responder,
    inputs=gr.Textbox(lines=2, label="Sua pergunta sobre Deep Learning"),
    outputs=gr.Textbox(lines=8, label="Resposta do Professor DL"),
    title="Professor DL - Bot de Deep Learning",
    description=(
        "Faça perguntas sobre redes neurais, Deep Learning, CNN, RNN, overfitting, "
        "regularização, etc. O Professor DL responde em português, de forma didática."
    ),
)

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