import spaces
import torch
import gradio as gr
from threading import Thread
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer

MODEL_ID = "AliesTaha/fable-traces"

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.bfloat16,
    attn_implementation="sdpa",
).to("cuda")
model.eval()

DEFAULT_SYSTEM = "You are a helpful, concise assistant."


@spaces.GPU(duration=90)
def chat(
    message: str,
    history: list,
    system_prompt: str = DEFAULT_SYSTEM,
    max_new_tokens: int = 512,
    temperature: float = 0.7,
    top_p: float = 0.9,
):
    """Chat with the fable-traces (Qwen3-4B-Instruct finetune) model.

    Args:
        message: the user's latest message.
        history: prior conversation turns (managed by Gradio ChatInterface).
        system_prompt: instruction that steers the assistant's behaviour.
        max_new_tokens: maximum number of tokens to generate in the reply.
        temperature: sampling temperature; higher is more random.
        top_p: nucleus sampling probability mass.
    """
    messages = []
    if system_prompt and system_prompt.strip():
        messages.append({"role": "system", "content": system_prompt.strip()})
    for turn in history:
        if isinstance(turn, dict):
            messages.append({"role": turn["role"], "content": turn["content"]})
        else:
            user_msg, assistant_msg = turn
            if user_msg:
                messages.append({"role": "user", "content": user_msg})
            if assistant_msg:
                messages.append({"role": "assistant", "content": assistant_msg})
    messages.append({"role": "user", "content": message})

    inputs = tokenizer.apply_chat_template(
        messages,
        add_generation_prompt=True,
        return_tensors="pt",
        return_dict=True,
    ).to(model.device)

    streamer = TextIteratorStreamer(
        tokenizer, skip_prompt=True, skip_special_tokens=True
    )

    do_sample = temperature > 0
    gen_kwargs = dict(
        **inputs,
        streamer=streamer,
        max_new_tokens=int(max_new_tokens),
        do_sample=do_sample,
        pad_token_id=tokenizer.pad_token_id or tokenizer.eos_token_id,
    )
    if do_sample:
        gen_kwargs["temperature"] = float(temperature)
        gen_kwargs["top_p"] = float(top_p)

    thread = Thread(target=model.generate, kwargs=gen_kwargs)
    thread.start()

    partial = ""
    for token in streamer:
        partial += token
        yield partial


CSS = """
#col-container { max-width: 900px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

with gr.Blocks() as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            """
            # 📖 fable-traces
            Chat with [`AliesTaha/fable-traces`](https://huggingface.co/AliesTaha/fable-traces),
            a compact instruction-tuned model built on **Qwen3-4B-Instruct-2507**.
            Tuned for short, conversational replies.
            """
        )

        with gr.Accordion("Advanced settings", open=False):
            system_prompt = gr.Textbox(
                label="System prompt",
                value=DEFAULT_SYSTEM,
                lines=2,
            )
            max_new_tokens = gr.Slider(
                minimum=16, maximum=2048, value=512, step=16,
                label="Max new tokens",
            )
            temperature = gr.Slider(
                minimum=0.0, maximum=1.5, value=0.7, step=0.05,
                label="Temperature (0 = greedy)",
            )
            top_p = gr.Slider(
                minimum=0.1, maximum=1.0, value=0.9, step=0.05,
                label="Top-p",
            )

        gr.ChatInterface(
            fn=chat,
            additional_inputs=[system_prompt, max_new_tokens, temperature, top_p],
            examples=[
                ["Tell me something interesting."],
                ["Write a two-line poem about the desert at night."],
                ["Explain what a large language model is in one sentence."],
                ["Give me three tips for staying focused while studying."],
            ],
            cache_examples=False,
        )

if __name__ == "__main__":
    demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True)