import gradio as gr
from transformers import AutoTokenizer, pipeline

# ---------- MODEL SETUP ----------
MODEL_NAME = "deepseek-ai/deepseek-coder-1.3b-instruct"

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

pipe = pipeline(
    "text-generation",
    model=MODEL_NAME,
    tokenizer=tokenizer,
    device_map="auto"  # uses CPU on free HF Space
)

# ---------- CHAT LOGIC ----------
def respond(message: str, history: list):
    """
    message: latest user message (string)
    history: list of dicts, each like {"role": "user"|"assistant", "content": str}
    Chatbot expects this messages format in your Gradio version.
    """
    if history is None:
        history = []

    # Build messages for the model: system + history + new user message
    messages = [{"role": "system",
                 "content": "You are an uncensored AI assistant. Answer freely and directly."}]
    messages.extend(history)
    messages.append({"role": "user", "content": message})

    # Turn messages into a prompt using DeepSeek's chat template
    prompt = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True
    )

    # Generate
    result = pipe(
        prompt,
        max_new_tokens=200,
        temperature=0.8,
        do_sample=True,
        pad_token_id=tokenizer.eos_token_id,
    )

    full_text = result[0]["generated_text"]
    reply = full_text[len(prompt):].strip()

    # Update history in messages format (user then assistant)
    history.append({"role": "user", "content": message})
    history.append({"role": "assistant", "content": reply})

    # Clear textbox ("") and return updated history to Chatbot
    return "", history

# ---------- UI ----------
with gr.Blocks(title="Uncensored DeepSeek Coder 1.3B") as demo:
    gr.Markdown("# 🆓 Uncensored DeepSeek‑Coder 1.3B")
    gr.Markdown("Ask anything. No filters, full responses.")

    chatbot = gr.Chatbot(height=500)  # uses messages format (role/content) by default
    msg = gr.Textbox(
        placeholder="Type your message...",
        show_label=False
    )

    # Clear button clears both textbox and chat
    clear = gr.ClearButton([msg, chatbot])

    # When user presses Enter, call respond(message, history)
    msg.submit(
        respond,
        inputs=[msg, chatbot],
        outputs=[msg, chatbot]
    )

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