Upload 6 files

README.md

@@ -1,13 +1,55 @@
 ---
-title: MicroMixer 2
-emoji: 🏆
-colorFrom: blue
-colorTo: gray
+title: MicroMixer-2 Discord Demo
+emoji: 🎛️
+colorFrom: purple
+colorTo: blue
 sdk: gradio
-sdk_version: 6.18.0
-python_version: '3.13'
+sdk_version: 4.44.0
 app_file: app.py
 pinned: false
+license: apache-2.0
+short_description: Play with the four attention-free MicroMixer-2 checkpoints.
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# MicroMixer-2 Discord-dialogue Playground
+
+Interactive Gradio demo for the four attention-free, MLP-only language
+models in [`llaa33219/micromixer-2`](https://huggingface.co/collections/llaa33219/micromixer-2).
+
+| Variant | Max context | Params |
+| --- | --- | --- |
+| 100K | 64  | ~125K |
+| 300K | 128 | ~431K |
+| 500K | 128 | ~779K |
+| 1M   | 4096| ~1.02M |
+
+All checkpoints are byte-level (vocabulary = 256) and were trained on
+[mookiezi/Discord-Dialogues](https://huggingface.co/datasets/mookiezi/Discord-Dialogues),
+so prompts in `User:` / `Assistant:` format tend to work best.
+
+## How it works
+
+- The first time you pick a model, `app.py` downloads its `model.pt` from
+  the Hub via `huggingface_hub.hf_hub_download` and caches it.
+- The architecture (`src/model.py`) and the byte-level tokenizer
+  (`src/tokenizer.py`) are vendored from
+  [llaa33219/MicroMixer-2](https://github.com/llaa33219/MicroMixer-2).
+- `MicroMixerConfig` is rebuilt from the exact hyperparameters listed on
+  each model card, so `load_state_dict` matches the published checkpoints.
+
+## Files
+
+```
+.
+├── app.py              # Gradio entry point
+├── requirements.txt
+├── src/
+│   ├── __init__.py
+│   ├── model.py        # MicroMixer-2 V4 architecture
+│   └── tokenizer.py    # ByteTokenizer (vocab=256)
+└── README.md
+```
+
+## License
+
+Apache 2.0, in line with the upstream project.

app.py

@@ -0,0 +1,316 @@
+"""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()

requirements.txt

@@ -0,0 +1,3 @@
+torch>=2.0
+gradio>=4.0
+huggingface_hub>=0.20

src/model.py

@@ -0,0 +1,526 @@
+"""MicroMixer-2 V4: MLP-Mixer architecture optimized for language models.
+
+V4 innovations based on research:
+- DropPath (stochastic depth): Regularization via random residual skipping
+- FourierMixing: Parameter-free FFT token mixing (FNet-inspired)
+- Padding-aware loss: Ignore padding tokens in cross-entropy
+- Label smoothing: Regularize overconfident predictions
+- Increased depth: 6-12 layers for larger models
+- HyperMixing (ACL 2023): O(S) token mixing via hypernetwork
+- RoPE: Rotary position embedding for length generalization
+- Standard MLP: Better knowledge capacity than GatedMLP
+"""
+
+import math
+from dataclasses import dataclass, field
+from enum import Enum, auto
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class TokenMixerType(Enum):
+    HYPER = auto()      # HyperMixing: O(S) via hypernetwork
+    FOURIER = auto()    # FourierMixing: O(S log S) via FFT, zero params
+
+
+class DropPath(nn.Module):
+    """Stochastic Depth (DropPath) per sample.
+    
+    Randomly drops entire residual branches during training.
+    Linear schedule: drop probability increases with layer depth.
+    """
+
+    def __init__(self, drop_prob: float = 0.0):
+        super().__init__()
+        self.drop_prob = drop_prob
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if not self.training or self.drop_prob == 0.0:
+            return x
+        
+        keep_prob = 1 - self.drop_prob
+        shape = (x.shape[0],) + (1,) * (x.ndim - 1)
+        random_tensor = torch.rand(shape, dtype=x.dtype, device=x.device)
+        random_tensor = torch.floor(random_tensor + keep_prob)
+        output = x / keep_prob * random_tensor
+        return output
+
+
+class MlpBlock(nn.Module):
+    """Standard 2-layer MLP with GELU activation."""
+
+    def __init__(self, in_dim: int, hidden_dim: int, dropout: float = 0.1):
+        super().__init__()
+        self.fc1 = nn.Linear(in_dim, hidden_dim)
+        self.fc2 = nn.Linear(hidden_dim, in_dim)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.fc1(x)
+        x = F.gelu(x)
+        x = self.dropout(x)
+        x = self.fc2(x)
+        x = self.dropout(x)
+        return x
+
+
+class RotaryPositionEmbedding(nn.Module):
+    """Rotary Position Embedding (RoPE) for length generalization."""
+
+    def __init__(self, dim: int, max_seq_len: int = 512):
+        super().__init__()
+        self.dim = dim
+        inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2).float() / dim))
+        self.register_buffer("inv_freq", inv_freq)
+
+    def forward(self, x: torch.Tensor, seq_len: int) -> torch.Tensor:
+        t = torch.arange(seq_len, device=x.device).type_as(self.inv_freq)
+        freqs = torch.einsum("i,j->ij", t, self.inv_freq)
+        emb = torch.cat((freqs, freqs), dim=-1)
+        cos, sin = emb.cos(), emb.sin()
+        cos = cos.unsqueeze(0)
+        sin = sin.unsqueeze(0)
+
+        x_rot = x[..., : self.dim]
+        x_rest = x[..., self.dim :] if x.shape[-1] > self.dim else None
+
+        x1, x2 = x_rot[..., ::2], x_rot[..., 1::2]
+        rotated = torch.stack([-x2, x1], dim=-1).flatten(-2)
+        x_rotated = x_rot * cos + rotated * sin
+
+        if x_rest is not None:
+            return torch.cat([x_rotated, x_rest], dim=-1)
+        return x_rotated
+
+
+class HyperMixing(nn.Module):
+    """HyperMixing: O(S) token mixing via cumulative-mean hypernetwork.
+    
+    Based on HyperMixer (ACL 2023). Uses running statistics to generate
+    mixing weights dynamically.
+    """
+
+    def __init__(self, hidden_dim: int, hyper_hidden_dim: int, dropout: float = 0.1):
+        super().__init__()
+        self.hidden_dim = hidden_dim
+        self.hyper = nn.Sequential(
+            nn.Linear(hidden_dim, hyper_hidden_dim),
+            nn.GELU(),
+            nn.Linear(hyper_hidden_dim, hidden_dim * 2),
+        )
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        B, S, H = x.shape
+        
+        # Cumulative mean for causal context
+        cumsum = torch.cumsum(x, dim=1)
+        counts = torch.arange(1, S + 1, device=x.device).view(1, S, 1).float()
+        pooled = cumsum / counts
+
+        # Hypernetwork generates affine transform weights
+        weights = self.hyper(pooled)
+        w1, w2 = weights.chunk(2, dim=-1)
+
+        # Affine mixing: scale + shift
+        x = x * w1 + w2
+        return self.dropout(x)
+
+
+class FourierMixing(nn.Module):
+    """FourierMixing: Parameter-free token mixing via FFT.
+    
+    Based on FNet (NAACL 2022). Replaces attention with 2D FFT.
+    - Zero learnable parameters for token mixing
+    - O(S log S) complexity
+    - 80% faster than attention on GPUs
+    
+    Causal property: FFT mixes all positions, so we apply
+    cumulative masking to maintain autoregressive property.
+    """
+
+    def __init__(self, hidden_dim: int, dropout: float = 0.1):
+        super().__init__()
+        self.dropout = nn.Dropout(dropout)
+        
+        # Learnable scaling for output (optional, helps stability)
+        self.scale = nn.Parameter(torch.ones(1) * 0.1)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        B, S, H = x.shape
+        
+        # Apply FFT along sequence dimension (dim=1)
+        # Real-valued FFT preserves real output
+        x_fft = torch.fft.fft(x, dim=1).real
+        
+        # Apply causal masking: each position only sees itself and prior
+        # Use cumulative sum to enforce causality
+        mask = torch.triu(torch.ones(S, S, device=x.device)).bool()
+        # More efficient: use cumulative mean like HyperMixing
+        cumsum = torch.cumsum(x_fft, dim=1)
+        counts = torch.arange(1, S + 1, device=x.device).view(1, S, 1).float()
+        x_causal = cumsum / counts
+        
+        # Blend original FFT output with causal version
+        x = x_fft * (1 - self.scale) + x_causal * self.scale
+        
+        return self.dropout(x)
+
+
+class MicroMixerLayer(nn.Module):
+    """Single MicroMixer layer with DropPath regularization.
+    
+    Architecture:
+    1. LayerNorm -> Token Mixing -> DropPath -> Residual
+    2. LayerNorm -> Channel Mixing (MLP) -> DropPath -> Residual
+    """
+
+    def __init__(
+        self,
+        hidden_dim: int,
+        hyper_hidden_dim: int,
+        channel_mlp_dim: int,
+        dropout: float = 0.1,
+        drop_path: float = 0.0,
+        mixer_type: TokenMixerType = TokenMixerType.HYPER,
+    ):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(hidden_dim)
+        self.norm2 = nn.LayerNorm(hidden_dim)
+        
+        if mixer_type == TokenMixerType.HYPER:
+            self.token_mixer = HyperMixing(hidden_dim, hyper_hidden_dim, dropout)
+        elif mixer_type == TokenMixerType.FOURIER:
+            self.token_mixer = FourierMixing(hidden_dim, dropout)
+        else:
+            raise ValueError(f"Unknown mixer type: {mixer_type}")
+        
+        self.channel_mlp = MlpBlock(hidden_dim, channel_mlp_dim, dropout)
+        self.drop_path = DropPath(drop_path) if drop_path > 0 else nn.Identity()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # Token mixing with stochastic depth
+        residual = x
+        x = self.norm1(x)
+        x = self.token_mixer(x)
+        x = residual + self.drop_path(x)
+
+        # Channel mixing with stochastic depth
+        residual = x
+        x = self.norm2(x)
+        x = self.channel_mlp(x)
+        x = residual + self.drop_path(x)
+
+        return x
+
+
+@dataclass
+class MicroMixerConfig:
+    """Configuration for MicroMixer-2 V4.
+    
+    Attributes:
+        vocab_size: Vocabulary size (256 for byte-level).
+        max_seq_len: Maximum sequence length.
+        hidden_dim: Hidden dimension for embeddings and mixer layers.
+        hyper_hidden_dim: Hidden dimension for HyperMixing hypernetwork.
+        channel_mlp_dim: Inner dimension of channel-mixing MLP.
+        num_layers: Number of mixer layers.
+        dropout: Dropout probability.
+        drop_path: DropPath probability (0 = disabled).
+        label_smoothing: Label smoothing for cross-entropy (0 = disabled).
+        tie_weights: Tie input/output embeddings.
+        mixer_type: Token mixing strategy (HYPER or FOURIER).
+        pad_token_id: Padding token ID for masked loss.
+    """
+
+    vocab_size: int = 256
+    max_seq_len: int = 128
+    hidden_dim: int = 128
+    hyper_hidden_dim: int = 64
+    channel_mlp_dim: int = 256
+    num_layers: int = 2
+    dropout: float = 0.1
+    drop_path: float = 0.0
+    label_smoothing: float = 0.0
+    tie_weights: bool = True
+    mixer_type: TokenMixerType = TokenMixerType.HYPER
+    pad_token_id: int = 0
+
+
+class MicroMixer(nn.Module):
+    """MicroMixer-2 V4: MLP-Mixer language model with research-backed innovations.
+    
+    V4 improvements:
+    - DropPath: Stochastic depth regularization
+    - FourierMixing: Optional parameter-free FFT mixing
+    - Padding-aware loss: Ignores padding tokens
+    - Label smoothing: Regularizes overconfident predictions
+    - Increased depth: Up to 12 layers for larger models
+    """
+
+    def __init__(self, config: MicroMixerConfig):
+        super().__init__()
+        self.config = config
+        self.vocab_size = config.vocab_size
+        self.max_seq_len = config.max_seq_len
+        self.hidden_dim = config.hidden_dim
+        self.pad_token_id = config.pad_token_id
+
+        self.token_embedding = nn.Embedding(config.vocab_size, config.hidden_dim)
+        self.rope = RotaryPositionEmbedding(config.hidden_dim, config.max_seq_len)
+        self.dropout = nn.Dropout(config.dropout)
+
+        # DropPath: linear schedule (increases with depth)
+        dpr = [x.item() for x in torch.linspace(0, config.drop_path, config.num_layers)]
+
+        self.mixer_layers = nn.ModuleList([
+            MicroMixerLayer(
+                config.hidden_dim,
+                config.hyper_hidden_dim,
+                config.channel_mlp_dim,
+                config.dropout,
+                drop_path=dpr[i],
+                mixer_type=config.mixer_type,
+            )
+            for i in range(config.num_layers)
+        ])
+
+        self.layer_norm = nn.LayerNorm(config.hidden_dim)
+        self.lm_head = nn.Linear(config.hidden_dim, config.vocab_size, bias=False)
+
+        if config.tie_weights:
+            self.lm_head.weight = self.token_embedding.weight
+
+        self.apply(self._init_weights)
+
+    def _init_weights(self, module: nn.Module):
+        if isinstance(module, nn.Linear):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+            if getattr(module, "bias", None) is not None:
+                torch.nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+        elif isinstance(module, nn.LayerNorm):
+            torch.nn.init.ones_(module.weight)
+            torch.nn.init.zeros_(module.bias)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        targets: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+        B, S = input_ids.shape
+
+        if S > self.max_seq_len:
+            input_ids = input_ids[:, -self.max_seq_len :]
+            S = self.max_seq_len
+            if targets is not None:
+                targets = targets[:, -self.max_seq_len :]
+            if attention_mask is not None:
+                attention_mask = attention_mask[:, -self.max_seq_len :]
+
+        token_emb = self.token_embedding(input_ids)
+        x = self.rope(token_emb, S)
+        x = self.dropout(x)
+
+        for layer in self.mixer_layers:
+            x = layer(x)
+
+        x = self.layer_norm(x)
+        logits = self.lm_head(x)
+
+        if targets is not None:
+            # Flatten for cross-entropy
+            logits_flat = logits.view(-1, self.vocab_size)
+            targets_flat = targets.view(-1)
+            
+            # Build ignore_mask: padding tokens AND positions after padding
+            if attention_mask is not None:
+                # Shift mask: predict token AFTER seeing context
+                # mask[i] = 1 means token i is real, so predicting i+1 is valid
+                shifted_mask = torch.ones_like(attention_mask)
+                shifted_mask[:, 1:] = attention_mask[:, :-1]
+                ignore_mask = (shifted_mask.view(-1) == 0)
+                pad_indices = ignore_mask.nonzero(as_tuple=True)[0]
+            else:
+                # No mask provided: only ignore explicit pad tokens in targets
+                pad_indices = (targets_flat == self.pad_token_id).nonzero(as_tuple=True)[0]
+            
+            # Compute loss with label smoothing and padding ignore
+            loss = F.cross_entropy(
+                logits_flat,
+                targets_flat,
+                ignore_index=self.pad_token_id if len(pad_indices) > 0 else -100,
+                label_smoothing=self.config.label_smoothing,
+            )
+            return logits, loss
+
+        return logits
+
+    @torch.no_grad()
+    def generate(
+        self,
+        input_ids: torch.Tensor,
+        max_new_tokens: int,
+        temperature: float = 1.0,
+        top_k: int | None = None,
+    ) -> torch.Tensor:
+        """Autoregressive text generation."""
+        self.eval()
+        device = next(self.parameters()).device
+        input_ids = input_ids.to(device)
+
+        for _ in range(max_new_tokens):
+            logits = self(input_ids)
+            logits = logits[:, -1, :]
+
+            if temperature == 0.0:
+                next_token = torch.argmax(logits, dim=-1, keepdim=True)
+            else:
+                logits = logits / temperature
+
+                if top_k is not None:
+                    v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
+                    logits[logits < v[:, [-1]]] = -float("Inf")
+
+                probs = F.softmax(logits, dim=-1)
+                next_token = torch.multinomial(probs, num_samples=1)
+
+            input_ids = torch.cat([input_ids, next_token], dim=1)
+
+        return input_ids
+
+
+def count_parameters(model: nn.Module) -> int:
+    """Count total trainable parameters."""
+    return sum(p.numel() for p in model.parameters() if p.requires_grad)
+
+
+
+
+def micromixer_100k() -> MicroMixerConfig:
+    """~100K parameter model for testing/experimentation."""
+    return MicroMixerConfig(
+        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,
+    )
+
+
+def micromixer_300k() -> MicroMixerConfig:
+    """~300K parameter model for small-scale experiments."""
+    return MicroMixerConfig(
+        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,
+    )
+
+
+def micromixer_500k() -> MicroMixerConfig:
+    """~500K parameter model for medium-scale experiments."""
+    return MicroMixerConfig(
+        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,
+    )
+
+
+def micromixer_1m() -> MicroMixerConfig:
+    """~1M parameter model for standard experiments."""
+    return MicroMixerConfig(
+        max_seq_len=256,
+        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,
+    )
+
+
+def micromixer_1m_long(max_seq_len: int = 4096) -> MicroMixerConfig:
+    """~1M parameter model with extended context length."""
+    return MicroMixerConfig(
+        max_seq_len=max_seq_len,
+        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,
+    )
+
+
+def micromixer_1m_fourier() -> MicroMixerConfig:
+    """~1M parameter model with FourierMixing (parameter-free token mixing)."""
+    return MicroMixerConfig(
+        max_seq_len=256,
+        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,
+        mixer_type=TokenMixerType.FOURIER,
+    )
+
+
+if __name__ == "__main__":
+    print("Testing MicroMixer-2 V4...")
+    
+    for name, config_fn in [
+        ("100k", micromixer_100k),
+        ("300k", micromixer_300k),
+        ("500k", micromixer_500k),
+        ("1M", micromixer_1m),
+        ("1M-fourier", micromixer_1m_fourier),
+    ]:
+        config = config_fn()
+        model = MicroMixer(config)
+        params = count_parameters(model)
+        print(f"  {name}: {params:,} parameters, {config.num_layers} layers")
+
+        batch_size = 2
+        seq_len = min(32, config.max_seq_len)
+        input_ids = torch.randint(0, config.vocab_size, (batch_size, seq_len))
+
+        logits = model(input_ids)
+        assert logits.shape == (batch_size, seq_len, config.vocab_size)
+
+        targets = torch.randint(0, config.vocab_size, (batch_size, seq_len))
+        logits, loss = model(input_ids, targets)
+        assert logits.shape == (batch_size, seq_len, config.vocab_size)
+        assert loss.dim() == 0
+
+        prompt = input_ids[:, :4]
+        gen_ids = model.generate(prompt, max_new_tokens=8, temperature=0.8, top_k=10)
+        assert gen_ids.shape == (batch_size, 12)
+
+        prefix = torch.randint(0, config.vocab_size, (1, 5))
+        extra = torch.randint(0, config.vocab_size, (1, 3))
+        logits_prefix = model(prefix)[:, -1, :]
+        logits_extended = model(torch.cat([prefix, extra], dim=1))[:, 4, :]
+        assert torch.allclose(logits_prefix, logits_extended, atol=1e-5)
+
+    print("All V4 tests passed!")

src/tokenizer.py

@@ -0,0 +1,116 @@
+"""Byte-level tokenizer for MicroMixer-1 language model."""
+
+
+class ByteTokenizer:
+    """Byte-level tokenizer with 256 vocabulary size.
+
+    Reserves byte values 0, 1, 2 for special tokens:
+    - 0: pad_token
+    - 1: bos_token (beginning of sequence)
+    - 2: eos_token (end of sequence)
+
+    All other byte values (3-255) represent raw UTF-8 bytes.
+    """
+
+    def __init__(self):
+        self.vocab_size = 256
+        # Special tokens (using byte values that are rare in text)
+        self.pad_token_id = 0
+        self.bos_token_id = 1  # Beginning of sequence
+        self.eos_token_id = 2  # End of sequence
+
+    def encode(self, text: str) -> list[int]:
+        """Encode string to list of byte IDs.
+
+        Args:
+            text: Input string to encode
+
+        Returns:
+            List of byte IDs with BOS token at start and EOS token at end
+        """
+        if not text:
+            # Empty string: just return BOS + EOS
+            return [self.bos_token_id, self.eos_token_id]
+
+        # Convert text to UTF-8 bytes
+        text_bytes = text.encode("utf-8")
+
+        # Map each byte to its integer value (0-255)
+        byte_ids = [b for b in text_bytes]
+
+        # Add BOS at start, EOS at end
+        return [self.bos_token_id] + byte_ids + [self.eos_token_id]
+
+    def decode(self, ids: list[int]) -> str:
+        """Decode list of byte IDs back to string.
+
+        Args:
+            ids: List of byte IDs
+
+        Returns:
+            Decoded string
+        """
+        if not ids:
+            return ""
+
+        # Filter out special tokens (pad, bos, eos)
+        byte_values = [
+            b for b in ids
+            if b not in (self.pad_token_id, self.bos_token_id, self.eos_token_id)
+        ]
+
+        if not byte_values:
+            return ""
+
+        # Convert byte values back to bytes object
+        byte_data = bytes(byte_values)
+
+        # Decode UTF-8 bytes to string, handle errors gracefully
+        return byte_data.decode("utf-8", errors="replace")
+
+    def encode_batch(
+        self, texts: list[str], max_length: int = None, padding: bool = True
+    ) -> dict:
+        """Encode a batch of texts with optional padding.
+
+        Args:
+            texts: List of strings to encode
+            max_length: Maximum sequence length (truncation if specified)
+            padding: Whether to pad sequences to longest in batch
+
+        Returns:
+            Dict with 'input_ids' (list of lists) and 'attention_mask' (list of lists)
+        """
+        # Encode each text
+        encoded = [self.encode(text) for text in texts]
+
+        # Get sequence lengths before padding/truncation
+        lengths = [len(ids) for ids in encoded]
+
+        # Truncate if max_length specified
+        if max_length is not None:
+            encoded = [ids[:max_length] for ids in encoded]
+
+        # Pad to longest sequence if padding=True
+        if padding and encoded:
+            max_seq_len = max(len(ids) for ids in encoded)
+            pad_token_id = self.pad_token_id
+
+            padded = []
+            attention_masks = []
+            for ids in encoded:
+                pad_len = max_seq_len - len(ids)
+                padded.append(ids + [pad_token_id] * pad_len)
+                attention_masks.append([1] * len(ids) + [0] * pad_len)
+
+            return {
+                "input_ids": padded,
+                "attention_mask": attention_masks,
+            }
+
+        # No padding
+        attention_masks = [[1] * len(ids) for ids in encoded]
+        return {
+            "input_ids": encoded,
+            "attention_mask": attention_masks,
+        }