import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM

# Nome do modelo no Hugging Face Hub
model_name = "wpbcpaz/fitness-chatbot-model"
#model_name = "openai-community/gpt2"


# Carregamento do tokenizer e modelo
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# Função de resposta com histórico
def responder(message, history):
    # Constrói o histórico como prompt
    prompt = ""
    for user_msg, bot_msg in history:
        prompt += f"Usuário: {user_msg}\nChatbot: {bot_msg}\n"
    prompt += f"Usuário: {message}\nChatbot:"

    # Tokenização e geração
    inputs = tokenizer(prompt, return_tensors="pt")
    outputs = model.generate(**inputs, max_new_tokens=150, pad_token_id=tokenizer.eos_token_id)
    resposta = tokenizer.decode(outputs[0], skip_special_tokens=True)

    # Extrai apenas a resposta nova
    resposta_final = resposta.split("Chatbot:")[-1].strip()
    return resposta_final

# Interface estilo chat
demo = gr.ChatInterface(
    responder,
    title="🦜 Papa Capim-7B Chatbot",
    description="Modelo de linguagem em português para instruções, explicações e geração de texto. Ideal para aplicações educacionais e assistentes virtuais.",
    theme="soft",
    examples=[
        "Explique o que é biomecânica.",
        "Quais são os benefícios da musculação para idosos?",
        "Monte um plano alimentar para quem quer ganhar massa muscular.",
        "Como posso melhorar minha postura durante o treino?"
    ],
)

# Lançamento da interface
demo.launch()