import gradio as gr
from transformers import pipeline

# Carrega modelo StarCoder otimizado para geração de código
generator = pipeline("text-generation", model="bigcode/starcoder", device=-1)

def chat(message, history):
    history = history or []
    response = generator(message, max_new_tokens=200, do_sample=True, temperature=0.2)[0]["generated_text"]
    # Pega apenas a resposta após a mensagem do usuário
    response = response[len(message):].strip()
    history.append((message, response))
    return history, history

with gr.Blocks(css="style.css") as demo:
    gr.Markdown("# 💬 StarCoder Chat
Escreva código, faça perguntas e receba respostas em **Markdown**.")

    chatbot = gr.Chatbot(elem_id="chatbot")
    msg = gr.Textbox(placeholder="Digite sua mensagem...")
    clear = gr.Button("Limpar")

    msg.submit(chat, [msg, chatbot], [chatbot, chatbot])
    clear.click(lambda: None, None, chatbot)

demo.launch()