app.py

"""Hugging Face Space app for the MicroMixer-2 Discord-dialogue model family.

The four checkpoints in the llaa33219/micromixer-2 collection are all
MLP-Mixer based byte-level language models. They differ only in size,
max sequence length, and a couple of regularisation knobs (DropPath,
label smoothing) - everything else shares the same architecture and
the same ByteTokenizer.

This app exposes a Gradio demo that:
  * lets the user pick one of the four checkpoints,
  * downloads `model.pt` from the Hub on first use and caches it,
  * generates Discord-style replies from a prompt.
"""

from __future__ import annotations

import os
from pathlib import Path
from typing import Dict, Tuple

import gradio as gr
import torch
from huggingface_hub import hf_hub_download

from src.model import MicroMixer, MicroMixerConfig
from src.tokenizer import ByteTokenizer


# ---------------------------------------------------------------------------
# Model registry
# ---------------------------------------------------------------------------
# Each entry maps a UI label to:
#   (hf_repo_id, MicroMixerConfig kwargs matching the card)
#
# Numbers were taken straight from the model cards on the Hub
# (llaa33219/MicroMixer-2-*-discord-dialogues).
MODEL_REGISTRY: Dict[str, Tuple[str, dict]] = {
    "100K  (max 64 tok,  ~125K params)": (
        "llaa33219/MicroMixer-2-100K-discord-dialogues",
        dict(
            max_seq_len=64,
            hidden_dim=84,
            hyper_hidden_dim=48,
            channel_mlp_dim=128,
            num_layers=3,
            dropout=0.1,
            drop_path=0.0,
            label_smoothing=0.0,
        ),
    ),
    "300K  (max 128 tok, ~431K params)": (
        "llaa33219/MicroMixer-2-300K-discord-dialogues",
        dict(
            max_seq_len=128,
            hidden_dim=128,
            hyper_hidden_dim=64,
            channel_mlp_dim=288,
            num_layers=4,
            dropout=0.1,
            drop_path=0.05,
            label_smoothing=0.05,
        ),
    ),
    "500K  (max 128 tok, ~779K params)": (
        "llaa33219/MicroMixer-2-500K-discord-dialogues",
        dict(
            max_seq_len=128,
            hidden_dim=176,
            hyper_hidden_dim=88,
            channel_mlp_dim=384,
            num_layers=4,
            dropout=0.1,
            drop_path=0.1,
            label_smoothing=0.05,
        ),
    ),
    "1M  (max 4096 tok, ~1.02M params)": (
        "llaa33219/MicroMixer-2-1M-discord-dialogues",
        dict(
            max_seq_len=4096,
            hidden_dim=168,
            hyper_hidden_dim=84,
            channel_mlp_dim=448,
            num_layers=5,
            dropout=0.1,
            drop_path=0.1,
            label_smoothing=0.1,
        ),
    ),
}


# ---------------------------------------------------------------------------
# Cached model loader
# ---------------------------------------------------------------------------
class ModelCache:
    """Lazily downloads, builds, and caches the four MicroMixer checkpoints."""

    def __init__(self) -> None:
        self._cache: Dict[str, Tuple[MicroMixer, "torch.device"]] = {}
        self._tokenizer = ByteTokenizer()
        # Prefer an explicit env override, then the Spaces persistent volume
        # (/data, only present when the Space opted in to persistent storage),
        # then huggingface_hub's default cache. Falling back gracefully means
        # the demo still works on a stock CPU Space.
        env_dir = os.environ.get("MICROMIXER_CACHE")
        if env_dir:
            self._cache_dir = Path(env_dir)
        elif Path("/data").is_dir():
            self._cache_dir = Path("/data")
        else:
            self._cache_dir = None  # let hf_hub_download use its default
        if self._cache_dir is not None:
            self._cache_dir.mkdir(parents=True, exist_ok=True)
        self._device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    @property
    def device(self) -> "torch.device":
        return self._device

    @property
    def tokenizer(self) -> ByteTokenizer:
        return self._tokenizer

    def get(self, label: str) -> MicroMixer:
        if label not in self._cache:
            self._cache[label] = self._load(label)
        return self._cache[label][0]

    def _load(self, label: str) -> Tuple[MicroMixer, "torch.device"]:
        repo_id, cfg_kwargs = MODEL_REGISTRY[label]
        config = MicroMixerConfig(**cfg_kwargs)

        # 1. Download the .pt file (cached locally between Space restarts).
        download_kwargs = {"repo_id": repo_id, "filename": "model.pt"}
        if self._cache_dir is not None:
            download_kwargs["cache_dir"] = str(self._cache_dir)
        ckpt_path = hf_hub_download(**download_kwargs)

        # 2. Build the model and load weights.
        model = MicroMixer(config)
        state = torch.load(ckpt_path, map_location=self._device, weights_only=False)
        # Checkpoints on the Hub store the weights under "model_state_dict";
        # be defensive in case a future upload drops the wrapper key.
        if isinstance(state, dict) and "model_state_dict" in state:
            state = state["model_state_dict"]
        model.load_state_dict(state)

        model.to(self._device)
        model.eval()
        return model, self._device


CACHE = ModelCache()


# ---------------------------------------------------------------------------
# Generation
# ---------------------------------------------------------------------------
def _resolve_max_new_tokens(prompt: str, max_seq_len: int) -> int:
    """Cap generation so the running window never exceeds the model's context."""
    prompt_len = len(CACHE.tokenizer.encode(prompt))
    # Keep at least a 1-token safety margin for the seed token.
    budget = max_seq_len - prompt_len - 1
    return max(1, budget)


def generate(
    model_label: str,
    prompt: str,
    max_new_tokens: int,
    temperature: float,
    top_k: int,
) -> str:
    if not prompt:
        return "⚠️  Prompt is empty - type something first."

    try:
        model = CACHE.get(model_label)
    except Exception as exc:  # pragma: no cover - surfaced to the UI
        return f"❌ Failed to load `{model_label}`:\n```\n{exc}\n```"

    cfg = MODEL_REGISTRY[model_label][1]
    max_seq_len = cfg["max_seq_len"]
    hard_cap = _resolve_max_new_tokens(prompt, max_seq_len)
    max_new_tokens = int(min(max_new_tokens, hard_cap))

    input_ids = CACHE.tokenizer.encode(prompt)
    input_tensor = torch.tensor([input_ids], dtype=torch.long, device=CACHE.device)

    with torch.no_grad():
        output_ids = model.generate(
            input_tensor,
            max_new_tokens=max_new_tokens,
            temperature=temperature,
            top_k=int(top_k),
        )

    full = CACHE.tokenizer.decode(output_ids[0].cpu().tolist())
    return full


# ---------------------------------------------------------------------------
# Gradio UI
# ---------------------------------------------------------------------------
DEFAULT_PROMPT = "User: Hello! How are you today?\nAssistant:"

EXAMPLES = [
    ["User: What games do you play?\nAssistant:"],
    ["User: Tell me a joke about programming.\nAssistant:"],
    ["User: I'm bored, what should I do?\nAssistant:"],
    ["User: Good morning!\nAssistant:"],
    ["User: Do you like pizza?\nAssistant:"],
]


def build_demo() -> gr.Blocks:
    # NOTE: keep kwarg names compatible with Gradio 4.x / 5.x / 6.x.
    # `theme` was moved from Blocks() to launch() in Gradio 6, so we
    # hand the theme to launch() and leave Blocks() vanilla.
    with gr.Blocks(
        title="MicroMixer-2 Discord Demo",
    ) as demo:
        gr.Markdown(
            """
            # 🎛️ MicroMixer-2 Discord-dialogue Playground

            Try the four attention-free, MLP-only language models from
            [`llaa33219/micromixer-2`](https://huggingface.co/collections/llaa33219/micromixer-2).
            All checkpoints are byte-level (vocab = 256) and were trained on
            [mookiezi/Discord-Dialogues](https://huggingface.co/datasets/mookiezi/Discord-Dialogues),
            so prompting with a `User:` / `Assistant:` turn works best.

            | Variant | Max context | Params |
            | --- | --- | --- |
            | 100K | 64  | ~125K |
            | 300K | 128 | ~431K |
            | 500K | 128 | ~779K |
            | 1M   | 4096| ~1.02M |
            """
        )

        with gr.Row():
            with gr.Column(scale=1):
                model_dd = gr.Dropdown(
                    choices=list(MODEL_REGISTRY.keys()),
                    value="1M  (max 4096 tok, ~1.02M params)",
                    label="Model",
                    info="The 1M model is the strongest but slowest.",
                )
                prompt_tb = gr.Textbox(
                    label="Prompt",
                    value=DEFAULT_PROMPT,
                    lines=4,
                    placeholder="User: ...\nAssistant:",
                )
                with gr.Accordion("Sampling settings", open=True):
                    max_new = gr.Slider(
                        minimum=8, maximum=512, value=128, step=8,
                        label="max_new_tokens",
                    )
                    temperature = gr.Slider(
                        minimum=0.1, maximum=2.0, value=0.8, step=0.05,
                        label="temperature",
                    )
                    top_k = gr.Slider(
                        minimum=0, maximum=200, value=40, step=1,
                        label="top_k (0 = off)",
                    )
                run_btn = gr.Button("Generate", variant="primary")
            with gr.Column(scale=2):
                output = gr.Textbox(
                    label="Output",
                    lines=18,
                    interactive=False,
                )

        gr.Examples(
            examples=EXAMPLES,
            inputs=[prompt_tb],
            label="Prompt examples (User/Assistant format)",
        )

        gr.Markdown(
            """
            ---
            ### Notes
            * First run for a given model downloads `model.pt` from the Hub
              (one-time, then cached). All four checkpoints together are < 20 MB.
            * The 100K/300K/500K models cap context at 64–128 bytes, so the
              UI clamps `max_new_tokens` automatically.
            * Runs on CPU by default; a CUDA GPU will be used automatically
              if the Space has one.
            * Source: [github.com/llaa33219/MicroMixer-2](https://github.com/llaa33219/MicroMixer-2)
            """
        )

        run_btn.click(
            fn=generate,
            inputs=[model_dd, prompt_tb, max_new, temperature, top_k],
            outputs=output,
        )

    return demo


if __name__ == "__main__":
    demo = build_demo()
    demo.queue(max_size=8).launch(
        server_name="0.0.0.0",
        server_port=7860,
        theme=gr.themes.Soft(),
    )
else:
    # When imported (e.g. by Spaces that wrap `app.py`).
    demo = build_demo()