import gradio as gr
from transformers import pipeline

# Nome do modelo no Hugging Face
MODEL_NAME = "DianoAI/Diano-Small"

# Pipeline do Hugging Face
pipe = pipeline("text2text-generation", model=MODEL_NAME)

# Função para gerar resposta
def chat(prompt):
    response = pipe(prompt, max_length=200)
    return [{"role": "assistant", "content": response[0]["generated_text"]}]

# Criar interface
with gr.Blocks(fill_height=True) as demo:
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("# DianoAI Chat")
            gr.Markdown(f"Este Space usa o modelo `{MODEL_NAME}`, servido pela API do Hugging Face.")
            gr.LoginButton("Fazer Login")

        with gr.Column(scale=2):
            chatbot = gr.Chatbot(type="messages")
            input_box = gr.Textbox(placeholder="Digite sua pergunta...")
            send_button = gr.Button("Enviar")

            send_button.click(chat, inputs=input_box, outputs=chatbot)

# Iniciar app
demo.launch()