"""Gradio demo for Distilled MoE Draft Speculative Decoding."""
import spaces
import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import time

# Model IDs
DENSE_MODEL = "bharadwajvyadavalli/qwen2.5-distilled-draft-fast"
BASE_MODEL = "Qwen/Qwen2.5-0.5B"

@spaces.GPU
def generate_comparison(prompt, max_tokens, temperature):
    """Generate text with both models and compare."""
    if not prompt.strip():
        return "", "0", "0 tok/s", "", "0", "0 tok/s"

    # Load tokenizer
    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)

    results = []

    for model_id in [DENSE_MODEL, BASE_MODEL]:
        model = AutoModelForCausalLM.from_pretrained(
            model_id,
            torch_dtype=torch.float16,
            device_map="auto"
        )

        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

        start_time = time.time()

        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=int(max_tokens),
                temperature=float(temperature) if temperature > 0 else 1.0,
                do_sample=temperature > 0,
                pad_token_id=tokenizer.eos_token_id,
            )

        elapsed = time.time() - start_time

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

        num_tokens = len(generated_ids)
        tokens_per_sec = num_tokens / elapsed if elapsed > 0 else 0

        results.extend([generated_text, str(num_tokens), f"{tokens_per_sec:.1f} tok/s"])

        # Free memory
        del model
        torch.cuda.empty_cache()

    return tuple(results)

# Build Gradio interface
with gr.Blocks(title="Distilled MoE Draft Demo") as demo:
    gr.Markdown("""
    # Distilled MoE Draft for Speculative Decoding

    Compare text generation between:
    - **Dense Distilled Draft**: KL-distilled from Qwen2.5
    - **Base Model**: Original Qwen2.5-0.5B

    [GitHub](https://github.com/bharadwajvyadavalli/distilled-moe-drafts) |
    [Dense Model](https://huggingface.co/bharadwajvyadavalli/qwen2.5-distilled-draft-fast) |
    [MoE Model](https://huggingface.co/bharadwajvyadavalli/qwen2.5-moe-draft-4experts-fast)
    """)

    with gr.Row():
        prompt = gr.Textbox(
            label="Prompt",
            placeholder="Enter your prompt here...",
            lines=3,
            value="Once upon a time in a land far away,"
        )

    with gr.Row():
        max_tokens = gr.Slider(10, 100, value=30, step=10, label="Max Tokens")
        temperature = gr.Slider(0, 1.5, value=0.7, step=0.1, label="Temperature")

    generate_btn = gr.Button("Generate", variant="primary")

    with gr.Row():
        with gr.Column():
            gr.Markdown("### Dense Distilled Draft")
            output1 = gr.Textbox(label="Generated Text", lines=4)
            with gr.Row():
                tokens1 = gr.Textbox(label="Tokens")
                speed1 = gr.Textbox(label="Speed")

        with gr.Column():
            gr.Markdown("### Base (Qwen2.5-0.5B)")
            output2 = gr.Textbox(label="Generated Text", lines=4)
            with gr.Row():
                tokens2 = gr.Textbox(label="Tokens")
                speed2 = gr.Textbox(label="Speed")

    generate_btn.click(
        generate_comparison,
        inputs=[prompt, max_tokens, temperature],
        outputs=[output1, tokens1, speed1, output2, tokens2, speed2]
    )

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