import spaces  # MUST come before torch / any CUDA-touching import
import torch
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer

MODEL_ID = "SupraLabs/supra-title-50m-pre"

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


@spaces.GPU(duration=10)
def generate_title(
    user_message: str,
    temperature: float = 0.4,
    top_p: float = 0.85,
    top_k: int = 40,
    repetition_penalty: float = 1.2,
    max_new_tokens: int = 24,
) -> str:
    """Generate a concise chat title from a user message.

    Args:
        user_message: The chat message or prompt to generate a title for.
        temperature: Sampling temperature — lower means more deterministic.
        top_p: Nucleus sampling probability threshold.
        top_k: Limits vocabulary to the top-K candidates.
        repetition_penalty: Penalizes repeated tokens.
        max_new_tokens: Maximum number of new tokens to generate.
    """
    prompt = f"User: {user_message}\nTitle: "
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            do_sample=True,
            temperature=temperature,
            top_p=top_p,
            top_k=top_k,
            repetition_penalty=repetition_penalty,
            pad_token_id=tokenizer.pad_token_id,
            eos_token_id=tokenizer.eos_token_id,
        )

    generated_ids = outputs[0][inputs["input_ids"].shape[1]:]
    title = tokenizer.decode(generated_ids, skip_special_tokens=True).strip()
    return title


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

with gr.Blocks(theme=gr.themes.Citrus(), css=CSS) as demo:
    gr.Markdown(
        "# Supra Title 50M — Chat Title Generator\n"
        "Generate concise 2–4 word chat titles from user messages using "
        "[SupraLabs/supra-title-50m-pre](https://huggingface.co/SupraLabs/supra-title-50m-pre), "
        "a lightweight 50M-parameter Llama model."
    )

    with gr.Column(elem_id="col-container"):
        with gr.Row():
            user_message = gr.Textbox(
                label="User message",
                placeholder="Type a chat message or prompt to title…",
                lines=2,
                scale=4,
            )
            run_btn = gr.Button("Generate Title", variant="primary", scale=1)

        title_output = gr.Textbox(label="Generated title", interactive=False)

        with gr.Accordion("Advanced settings", open=False):
            temperature = gr.Slider(0.1, 1.5, value=0.4, step=0.05, label="Temperature")
            top_p = gr.Slider(0.1, 1.0, value=0.85, step=0.05, label="Top-p")
            top_k = gr.Slider(1, 100, value=40, step=1, label="Top-k")
            repetition_penalty = gr.Slider(
                1.0, 2.0, value=1.2, step=0.05, label="Repetition penalty"
            )
            max_new_tokens = gr.Slider(
                4, 64, value=24, step=1, label="Max new tokens"
            )

    examples = [
        "How does AI work?",
        "Who was Albert Einstein?",
        "How to make a discord server?",
        "Why does my WiFi keep disconnecting the whole time??",
        "How to bake a cake?",
        "Why is the sky blue?",
        "How does a car work?",
        "What is Google?",
        "How does Google Maps work?",
        "Who is Elon Musk?",
        "Name the top three products of Microsoft",
    ]

    gr.Examples(
        examples=[[e] for e in examples],
        inputs=[user_message],
        outputs=title_output,
        fn=generate_title,
        cache_examples=True,
        cache_mode="lazy",
    )

    run_btn.click(
        fn=generate_title,
        inputs=[user_message, temperature, top_p, top_k, repetition_penalty, max_new_tokens],
        outputs=title_output,
        api_name="generate_title",
    )
    user_message.submit(
        fn=generate_title,
        inputs=[user_message, temperature, top_p, top_k, repetition_penalty, max_new_tokens],
        outputs=title_output,
        api_name="generate_title_submit",
    )

demo.launch(mcp_server=True)