"""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()