Upload 88 files

downloads/CompactAI Studio Setup 1.0.0.exe

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:245e455f3dd33a7d2d1f91f456d9aa926705b7bd69c78f068e8c56e3e846ae3a
-size 133
+oid sha256:f685d561cb3d1e9b8f41bfea7b50c8d5bd0b72007000c7f70f63747127c5a57f
+size 128

downloads/CompactAI Studio-1.0.0.AppImage

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:82ac5097ddba3db79ec3a05998820f44e1cd536ce60ae1eeaae08b391b466d1d
-size 133
+oid sha256:fa6103e15bfdc80bfea40471ef418e54ca4f6f8b6c90ec166072d821d93dfe3c
+size 128

downloads/interactive.py

@@ -18,6 +18,8 @@ from typing import Any, Dict, Iterator, List, Optional, Sequence, Tuple
 from urllib.parse import quote, unquote, urlparse
 from urllib.request import Request, urlopen
 
+import hashlib
+
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -41,10 +43,17 @@ class ModelConfig:
     seq_len: int = 2048
     sliding_window_size: int = 512
     mtp_horizons: Tuple[int, ...] = (2, 3, 4)
-    rope_fraction: float = 0.25
+    rope_fraction: float = 0.5
     embed_scale: bool = True
     logit_soft_cap: float = -1.0
     quantization: str = "nvfp4"
+    # Engram (conditional memory) config
+    engram_dim: int = 0
+    engram_heads: int = 4
+    engram_table_size: int = 8192
+    engram_max_ngram: int = 3
+    # mHC (Manifold-Constrained Hyper-Connections) config
+    mhc_expansion: int = 1
 
     @property
     def head_dim(self) -> int:
@@ -55,16 +64,17 @@ model_config = ModelConfig()
 
 MODEL_SERIES = {
     "haiku": {
-        "dim": 128,
-        "n_unique_layers": 8,
-        "n_logical_layers": 16,
+        "dim": 64,
+        "n_unique_layers": 12,
+        "n_logical_layers": 24,
         "n_heads": 4,
         "n_kv_heads": 2,
-        "ffn_dim": 224,
+        "ffn_dim": 384,
         "dropout": 0.0,
         "seq_len": 2048,
         "mtp_horizons": (2, 3, 4),
-        "batch_size": 48,
+        "rope_fraction": 0.5,
+        "batch_size": 80,
         "grad_accum": 1,
         "lr": 8e-4,
         "min_lr": 1e-5,
@@ -74,29 +84,34 @@ MODEL_SERIES = {
         "weight_decay": 0.02,
         "pretrain_passes": 2,
         "sft_passes": 3,
-        "max_sft_target_chars": 128,
-        "use_grad_checkpoint": False,
-        "use_torch_compile": True,
+        "max_sft_target_chars": 0,
+        "use_grad_checkpoint": True,
         "num_workers": 24,
         "prefetch_factor": 64,
         "shuffle_buffer": 8192,
         "max_pretrain_tokens": 0,
         "min_pretrain_tokens": 100_000_000,
         "quantization": "nvfp4",
+        "engram_dim": 8,
+        "engram_heads": 2,
+        "engram_table_size": 64,
+        "engram_max_ngram": 2,
+        "mhc_expansion": 2,
     },
     "sonnet": {
-        "dim": 768,
-        "n_unique_layers": 18,
-        "n_logical_layers": 36,
-        "n_heads": 12,
+        "dim": 1024,
+        "n_unique_layers": 20,
+        "n_logical_layers": 40,
+        "n_heads": 16,
         "n_kv_heads": 4,
-        "ffn_dim": 2538,
+        "ffn_dim": 4096,
         "dropout": 0.0,
         "seq_len": 2048,
         "mtp_horizons": (2,),
-        "batch_size": 6,
+        "rope_fraction": 0.5,
+        "batch_size": 24,
         "grad_accum": 1,
-        "lr": 2e-4,
+        "lr": 1e-4,
         "min_lr": 2e-5,
         "sft_lr": 5e-5,
         "sft_min_lr": 5e-6,
@@ -104,27 +119,32 @@ MODEL_SERIES = {
         "weight_decay": 0.1,
         "pretrain_passes": 1,
         "sft_passes": 1,
-        "max_sft_target_chars": 512,
+        "max_sft_target_chars": 0,
         "use_grad_checkpoint": True,
-        "use_torch_compile": True,
         "num_workers": 32,
-        "prefetch_factor": 48,
+        "prefetch_factor": 64,
         "shuffle_buffer": 16384,
         "max_pretrain_tokens": 0,
-        "min_pretrain_tokens": 0,
+        "min_pretrain_tokens": 100_000_000,
         "quantization": "nvfp4",
+        "engram_dim": 32,
+        "engram_heads": 8,
+        "engram_table_size": 4096,
+        "engram_max_ngram": 2,
+        "mhc_expansion": 2,
     },
     "opus": {
-        "dim": 1024,
-        "n_unique_layers": 20,
-        "n_logical_layers": 40,
+        "dim": 1536,
+        "n_unique_layers": 18,
+        "n_logical_layers": 36,
         "n_heads": 16,
         "n_kv_heads": 4,
-        "ffn_dim": 3557,
+        "ffn_dim": 5888,
         "dropout": 0.0,
         "seq_len": 2048,
         "mtp_horizons": (2,),
-        "batch_size": 12,
+        "rope_fraction": 0.5,
+        "batch_size": 24,
         "grad_accum": 1,
         "lr": 1.6e-4,
         "min_lr": 1.6e-5,
@@ -134,15 +154,19 @@ MODEL_SERIES = {
         "weight_decay": 0.1,
         "pretrain_passes": 1,
         "sft_passes": 1,
-        "max_sft_target_chars": 1024,
+        "max_sft_target_chars": 0,
         "use_grad_checkpoint": True,
-        "use_torch_compile": True,
         "num_workers": 48,
-        "prefetch_factor": 48,
+        "prefetch_factor": 64,
         "shuffle_buffer": 16384,
         "max_pretrain_tokens": 0,
-        "min_pretrain_tokens": 0,
+        "min_pretrain_tokens": 100_000_000,
         "quantization": "nvfp4",
+        "engram_dim": 64,
+        "engram_heads": 8,
+        "engram_table_size": 8192,
+        "engram_max_ngram": 2,
+        "mhc_expansion": 4,
     },
 }
 
@@ -381,6 +405,10 @@ class CausalSelfAttention(nn.Module):
         self.wv = nn.Linear(dim, n_kv_heads * head_dim, bias=False)
         self.wo = nn.Linear(n_heads * head_dim, dim, bias=False)
 
+        for lin in (self.wq, self.wk, self.wv):
+            nn.init.normal_(lin.weight, std=dim ** -0.5)
+        nn.init.normal_(self.wo.weight, std=(n_heads * head_dim) ** -0.5)
+
         self.rope_dim = max(2, int(head_dim * rope_fraction) // 2 * 2)
         self.rope = RotaryEmbedding(self.rope_dim)
 
@@ -444,7 +472,7 @@ class CausalSelfAttention(nn.Module):
                 .reshape(B, self.n_heads, S, self.head_dim)
             )
 
-        drop_p = self.dropout if self.training else 0.0
+        drop_p = self.dropout if (self.training and torch.is_grad_enabled()) else 0.0
 
         if is_global:
             if past_kv is None and T > 1:
@@ -479,9 +507,184 @@ class SwiGLU(nn.Module):
         self.down = nn.Linear(hidden_dim, dim, bias=False)
         self.drop = nn.Dropout(dropout)
 
+        nn.init.normal_(self.gate.weight, std=dim ** -0.5)
+        nn.init.normal_(self.up.weight, std=dim ** -0.5)
+        nn.init.normal_(self.down.weight, std=hidden_dim ** -0.5)
+
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         h = F.silu(self.gate(x)) * self.up(x)
-        return self.drop(self.down(h))
+        out = self.down(h)
+        if self.training and torch.is_grad_enabled():
+            out = self.drop(out)
+        return out
+
+
+class EngramBlock(nn.Module):
+    """Conditional memory via O(1) hashed N-gram lookup (DeepSeek Engram)."""
+
+    def __init__(
+        self,
+        dim: int,
+        engram_dim: int,
+        n_heads: int = 4,
+        table_size: int = 8192,
+        max_ngram: int = 3,
+    ) -> None:
+        super().__init__()
+        self.dim = dim
+        self.engram_dim = engram_dim
+        self.n_heads = n_heads
+        self.table_size = table_size
+        self.max_ngram = max_ngram
+
+        self.embeddings = nn.ParameterDict()
+        for n in range(2, max_ngram + 1):
+            for k in range(n_heads):
+                self.embeddings[f"{n}_{k}"] = nn.Parameter(
+                    torch.randn(table_size, engram_dim) * (engram_dim**-0.5)
+                )
+
+        for n in range(2, max_ngram + 1):
+            for k in range(n_heads):
+                seed = int(hashlib.md5(f"engram_{n}_{k}".encode()).hexdigest()[:8], 16)
+                rng = torch.Generator().manual_seed(seed)
+                a = torch.randint(1, 2**31, (1,), generator=rng).item()
+                b = torch.randint(0, 2**31, (1,), generator=rng).item()
+                self.register_buffer(
+                    f"hash_a_{n}_{k}", torch.tensor(a), persistent=False
+                )
+                self.register_buffer(
+                    f"hash_b_{n}_{k}", torch.tensor(b), persistent=False
+                )
+
+        total_branch_dim = engram_dim * n_heads * (max_ngram - 1)
+        self.branch_conv = nn.Conv1d(
+            total_branch_dim,
+            total_branch_dim,
+            kernel_size=4,
+            dilation=max_ngram,
+            padding=0,
+            groups=total_branch_dim,
+            bias=True,
+        )
+        nn.init.zeros_(self.branch_conv.weight)
+        nn.init.zeros_(self.branch_conv.bias)
+
+        self.gate_query = nn.Linear(dim, engram_dim, bias=False)
+        self.gate_key = nn.Linear(total_branch_dim, engram_dim, bias=False)
+        self.gate_value = nn.Linear(total_branch_dim, dim, bias=False)
+        self.gate_scale = engram_dim**-0.5
+
+    def _hash_ngram(self, token_ids: torch.Tensor, n: int, k: int) -> torch.Tensor:
+        a = getattr(self, f"hash_a_{n}_{k}")
+        b = getattr(self, f"hash_b_{n}_{k}")
+        B, T = token_ids.shape
+        padded = F.pad(token_ids, (n - 1, 0), value=0)
+        combined = torch.zeros(B, T, dtype=torch.long, device=token_ids.device)
+        for i in range(n):
+            combined = (combined * 31 + padded[:, i : i + T].long()) % self.table_size
+        return ((a * combined) ^ b) % self.table_size
+
+    def forward(
+        self, hidden: torch.Tensor, token_ids: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        B, T, _ = hidden.shape
+        if token_ids is None:
+            token_ids = hidden.mean(dim=-1).long() % self.table_size
+        all_indices = []
+        all_tables = []
+        for n in range(2, self.max_ngram + 1):
+            for k in range(self.n_heads):
+                all_indices.append(self._hash_ngram(token_ids, n, k))
+                all_tables.append(self.embeddings[f"{n}_{k}"])
+        branch_outputs = [tbl[idx] for idx, tbl in zip(all_indices, all_tables)]
+        memory = torch.cat(branch_outputs, dim=-1)
+        conv_in = memory.transpose(1, 2)
+        conv_in = F.pad(
+            conv_in,
+            (self.branch_conv.dilation[0] * (self.branch_conv.kernel_size[0] - 1), 0),
+        )
+        conv_out = self.branch_conv(conv_in)
+        memory = conv_out.transpose(1, 2)
+        query = self.gate_query(hidden)
+        key = self.gate_key(memory)
+        gate = torch.sigmoid(
+            (query * key).sum(dim=-1, keepdim=True) * self.gate_scale
+        )
+        value = self.gate_value(memory)
+        return gate * value
+
+
+def _sinkhorn_knopp(logits: torch.Tensor, n_iters: int = 7) -> torch.Tensor:
+    M = torch.exp(logits.clamp(-10, 10))
+    for _ in range(n_iters):
+        M = M / M.sum(dim=-1, keepdim=True).clamp(min=1e-10)
+        M = M / M.sum(dim=-2, keepdim=True).clamp(min=1e-10)
+    return M
+
+
+class ManifoldHyperConnection(nn.Module):
+    """Manifold-Constrained Hyper-Connections (mHC) residual wrapper."""
+
+    def __init__(self, dim: int, expansion: int = 2) -> None:
+        super().__init__()
+        self.dim = dim
+        self.expansion = expansion
+        n = expansion
+
+        self.bias_pre = nn.Parameter(torch.zeros(1, n))
+        self.bias_post = nn.Parameter(torch.zeros(1, n))
+        self.bias_res = nn.Parameter(torch.zeros(n, n))
+
+        self.theta_pre = nn.Linear(n * dim, n, bias=False)
+        self.theta_post = nn.Linear(n * dim, n, bias=False)
+        self.theta_res = nn.Linear(n * dim, n * n, bias=False)
+
+        self.alpha_pre = nn.Parameter(torch.tensor(0.0))
+        self.alpha_post = nn.Parameter(torch.tensor(0.0))
+        self.alpha_res = nn.Parameter(torch.tensor(0.0))
+
+    def _compute_mappings(
+        self, x_expanded: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        B, T, _ = x_expanded.shape
+        n = self.expansion
+        x_norm = F.rms_norm(x_expanded, [x_expanded.shape[-1]])
+        d_pre = torch.tanh(self.theta_pre(x_norm))
+        d_post = torch.tanh(self.theta_post(x_norm))
+        d_res = self.theta_res(x_norm)
+        H_pre_raw = torch.sigmoid(self.alpha_pre * d_pre + self.bias_pre)
+        H_post_raw = 2.0 * torch.sigmoid(self.alpha_post * d_post + self.bias_post)
+        H_res_raw = (self.alpha_res * d_res + self.bias_res.reshape(1, 1, -1)).reshape(
+            B, T, n, n
+        )
+        H_res = _sinkhorn_knopp(H_res_raw)
+        return H_pre_raw.unsqueeze(-2), H_post_raw.unsqueeze(-2), H_res
+
+    def expand_stream(self, x: torch.Tensor) -> torch.Tensor:
+        return x.repeat(1, 1, self.expansion)
+
+    def collapse_stream(self, x_expanded: torch.Tensor) -> torch.Tensor:
+        B, T, _ = x_expanded.shape
+        return x_expanded.view(B, T, self.expansion, self.dim).mean(dim=-2)
+
+    def pre_mix(self, x_expanded: torch.Tensor, H_pre: torch.Tensor) -> torch.Tensor:
+        B, T, _ = x_expanded.shape
+        x_streams = x_expanded.view(B, T, self.expansion, self.dim)
+        return (H_pre @ x_streams).squeeze(-2)
+
+    def post_res_mix(
+        self,
+        layer_output: torch.Tensor,
+        x_expanded: torch.Tensor,
+        H_post: torch.Tensor,
+        H_res: torch.Tensor,
+    ) -> torch.Tensor:
+        B, T, _ = x_expanded.shape
+        x_streams = x_expanded.view(B, T, self.expansion, self.dim)
+        mixed = torch.matmul(H_res, x_streams)
+        post_out = torch.matmul(H_post.transpose(-2, -1), layer_output.unsqueeze(-2))
+        return (mixed + post_out).reshape(B, T, self.expansion * self.dim)
 
 
 class TransformerBlock(nn.Module):
@@ -495,8 +698,14 @@ class TransformerBlock(nn.Module):
         dropout: float,
         sliding_window: int,
         rope_fraction: float,
+        engram_dim: int = 0,
+        engram_heads: int = 4,
+        engram_table_size: int = 8192,
+        engram_max_ngram: int = 3,
+        mhc_expansion: int = 1,
     ) -> None:
         super().__init__()
+        self.dim = dim
         self.norm1 = RMSNorm(dim)
         self.attn = CausalSelfAttention(
             dim=dim,
@@ -509,6 +718,20 @@ class TransformerBlock(nn.Module):
         )
         self.norm2 = RMSNorm(dim)
         self.ffn = SwiGLU(dim, ffn_dim, dropout)
+        self.use_engram = engram_dim > 0
+        if self.use_engram:
+            self.engram = EngramBlock(
+                dim=dim,
+                engram_dim=engram_dim,
+                n_heads=engram_heads,
+                table_size=engram_table_size,
+                max_ngram=engram_max_ngram,
+            )
+            self.engram_norm = RMSNorm(dim)
+        self.use_mhc = mhc_expansion > 1
+        if self.use_mhc:
+            self.mhc_attn = ManifoldHyperConnection(dim, expansion=mhc_expansion)
+            self.mhc_ffn = ManifoldHyperConnection(dim, expansion=mhc_expansion)
 
     def forward(
         self,
@@ -516,13 +739,50 @@ class TransformerBlock(nn.Module):
         is_global: bool,
         past_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
         use_cache: bool = False,
+        token_ids: Optional[torch.Tensor] = None,
     ) -> Tuple[torch.Tensor, Optional[Tuple[torch.Tensor, torch.Tensor]]]:
-        attn_out, new_kv = self.attn(self.norm1(x), is_global, past_kv, use_cache)
-        x = x + attn_out
-        x = x + self.ffn(self.norm2(x))
+        if self.use_mhc:
+            x_exp = self.mhc_attn.expand_stream(x)
+            H_pre, H_post, H_res = self.mhc_attn._compute_mappings(x_exp)
+            attn_in = self.mhc_attn.pre_mix(x_exp, H_pre)
+            attn_out, new_kv = self.attn(
+                self.norm1(attn_in), is_global, past_kv, use_cache
+            )
+            x_exp = self.mhc_attn.post_res_mix(attn_out, x_exp, H_post, H_res)
+            if self.use_engram:
+                collapsed = self.mhc_attn.collapse_stream(x_exp)
+                collapsed = collapsed + self.engram(
+                    self.engram_norm(collapsed), token_ids=token_ids
+                )
+                x_exp = self.mhc_attn.expand_stream(collapsed)
+            H_pre2, H_post2, H_res2 = self.mhc_ffn._compute_mappings(x_exp)
+            ffn_in = self.mhc_ffn.pre_mix(x_exp, H_pre2)
+            ffn_out = self.ffn(self.norm2(ffn_in))
+            x_exp = self.mhc_ffn.post_res_mix(ffn_out, x_exp, H_post2, H_res2)
+            x = self.mhc_attn.collapse_stream(x_exp)
+        else:
+            attn_out, new_kv = self.attn(self.norm1(x), is_global, past_kv, use_cache)
+            x = x + attn_out
+            if self.use_engram:
+                x = x + self.engram(self.engram_norm(x), token_ids=token_ids)
+            x = x + self.ffn(self.norm2(x))
         return x, new_kv
 
 
+def _detect_engram_dim(state_dict: dict) -> int:
+    for key in state_dict:
+        if ".engram." in key and ".embeddings." in key:
+            return state_dict[key].shape[-1]
+    return 0
+
+
+def _detect_mhc_expansion(state_dict: dict) -> int:
+    for key, val in state_dict.items():
+        if ".mhc_attn.bias_pre" in key and val.dim() == 2:
+            return val.shape[-1]
+    return 1
+
+
 class TinyMemoryLM(nn.Module):
     def __init__(
         self,
@@ -537,8 +797,13 @@ class TinyMemoryLM(nn.Module):
         mtp_horizons: Sequence[int],
         grad_checkpoint: bool,
         sliding_window: int = 512,
-        rope_fraction: float = 0.25,
+        rope_fraction: float = 0.5,
         embed_scale: bool = True,
+        engram_dim: int = 0,
+        engram_heads: int = 4,
+        engram_table_size: int = 8192,
+        engram_max_ngram: int = 3,
+        mhc_expansion: int = 1,
     ) -> None:
         super().__init__()
         self.dim = dim
@@ -565,6 +830,11 @@ class TinyMemoryLM(nn.Module):
                     dropout=dropout,
                     sliding_window=sliding_window,
                     rope_fraction=rope_fraction,
+                    engram_dim=engram_dim,
+                    engram_heads=engram_heads,
+                    engram_table_size=engram_table_size,
+                    engram_max_ngram=engram_max_ngram,
+                    mhc_expansion=mhc_expansion,
                 )
                 for _ in range(n_unique_layers)
             ]
@@ -634,7 +904,7 @@ class TinyMemoryLM(nn.Module):
         )
 
         for layer_idx, (block, logical_idx) in enumerate(logical_layers):
-            is_global = logical_idx % 3 == 0
+            is_global = logical_idx % 2 == 0
             past_kv = (
                 past_key_values[layer_idx]
                 if past_key_values is not None and layer_idx < len(past_key_values)
@@ -643,17 +913,20 @@ class TinyMemoryLM(nn.Module):
 
             if self.grad_checkpoint and self.training and not use_cache:
                 x, layer_kv = checkpoint(
-                    block, x, is_global, past_kv, use_cache, use_reentrant=False
+                    block, x, is_global, past_kv, use_cache, ids, use_reentrant=True
                 )
             else:
-                x, layer_kv = block(x, is_global, past_kv, use_cache)
+                x, layer_kv = block(x, is_global, past_kv, use_cache, ids)
 
             if new_past_key_values is not None:
                 new_past_key_values.append(layer_kv)
 
         x = self.norm(x)
         h_out = x if return_hidden else None
-        logits = self.head(x) + self.output_bias
+        logits = self.head(x)
+        if self.embed_scale_factor != 1.0:
+            logits = logits / self.embed_scale_factor
+        logits = logits + self.output_bias
 
         mtp: Dict[int, torch.Tensor] = {}
         if self.mtp_horizons and self.training:
@@ -662,7 +935,10 @@ class TinyMemoryLM(nn.Module):
                     shifted_h = x[:, :-horizon, :]
                     adapted_h = self.mtp_adapters[str(horizon)](shifted_h)
                     adapted_h = self.mtp_norms[str(horizon)](adapted_h)
-                    mtp_logits = self.head(adapted_h) + self.output_bias
+                    mtp_logits = self.head(adapted_h)
+                    if self.embed_scale_factor != 1.0:
+                        mtp_logits = mtp_logits / self.embed_scale_factor
+                    mtp_logits = mtp_logits + self.output_bias
                     mtp[horizon] = mtp_logits
 
         return logits, mtp, h_out, new_past_key_values
@@ -1462,6 +1738,14 @@ def load_bundle(repo_id: str, model_type: str) -> dict[str, object]:
         ckpt = torch.load(str(model_path), map_location="cpu", weights_only=False)
         cfg = series_config(series)
         vocab_size = int(ckpt.get("vocab_size", tokenizer.vocab_size))
+        state_dict = ckpt.get("model_state") or ckpt.get("state_dict") or ckpt
+        # Auto-detect new arch features from checkpoint weights
+        engram_dim = _detect_engram_dim(state_dict) or int(
+            cfg.get("engram_dim", model_config.engram_dim)
+        )
+        mhc_expansion = _detect_mhc_expansion(state_dict) or int(
+            cfg.get("mhc_expansion", model_config.mhc_expansion)
+        )
         model = TinyMemoryLM(
             vocab_size=vocab_size,
             dim=int(cfg.get("dim", model_config.dim)),
@@ -1480,19 +1764,24 @@ def load_bundle(repo_id: str, model_type: str) -> dict[str, object]:
             ),
             grad_checkpoint=False,
             sliding_window=int(
-                cfg.get(
-                    "sliding_window_size",
-                    getattr(model_config, "sliding_window_size", 512),
-                )
+                cfg.get("sliding_window_size", model_config.sliding_window_size)
             ),
             rope_fraction=float(
-                cfg.get("rope_fraction", getattr(model_config, "rope_fraction", 0.25))
+                cfg.get("rope_fraction", model_config.rope_fraction)
             ),
             embed_scale=bool(
-                cfg.get("embed_scale", getattr(model_config, "embed_scale", True))
+                cfg.get("embed_scale", model_config.embed_scale)
+            ),
+            engram_dim=engram_dim,
+            engram_heads=int(cfg.get("engram_heads", model_config.engram_heads)),
+            engram_table_size=int(
+                cfg.get("engram_table_size", model_config.engram_table_size)
             ),
+            engram_max_ngram=int(
+                cfg.get("engram_max_ngram", model_config.engram_max_ngram)
+            ),
+            mhc_expansion=mhc_expansion,
         )
-        state_dict = ckpt.get("model_state") or ckpt.get("state_dict") or ckpt
         model.load_state_dict(state_dict, strict=False)
         model.eval()
         if tokenizer.vocab_size > vocab_size:
@@ -1678,7 +1967,7 @@ def page_html() -> str:
       </div>
       <div class="meta">
         <span class="chip">Hugging Face: CompactAI</span>
-        <span class="chip">Auto-installs deps</span>
+        <span class="chip">pip install -r requirements.txt</span>
         <span class="chip">Local inference</span>
       </div>
     </div>

interactive.py

@@ -0,0 +1,2277 @@
+#!/usr/bin/env python3
+from __future__ import annotations
+
+import json
+import math
+import os
+import re
+import shutil
+import socket
+import string
+import sys
+import threading
+import webbrowser
+from dataclasses import dataclass
+from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
+from pathlib import Path
+from typing import Any, Dict, Iterator, List, Optional, Sequence, Tuple
+from urllib.parse import quote, unquote, urlparse
+from urllib.request import Request, urlopen
+
+import hashlib
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.checkpoint import checkpoint
+
+
+# ---------------------------------------------------------------------------
+# Config (from ailay.config)
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class ModelConfig:
+    dim: int = 128
+    n_unique_layers: int = 8
+    n_logical_layers: int = 16
+    n_heads: int = 4
+    n_kv_heads: int = 2
+    ffn_dim: int = 224
+    dropout: float = 0.0
+    seq_len: int = 2048
+    sliding_window_size: int = 512
+    mtp_horizons: Tuple[int, ...] = (2, 3, 4)
+    rope_fraction: float = 0.5
+    embed_scale: bool = True
+    logit_soft_cap: float = -1.0
+    quantization: str = "nvfp4"
+    # Engram (conditional memory) config
+    engram_dim: int = 0
+    engram_heads: int = 4
+    engram_table_size: int = 8192
+    engram_max_ngram: int = 3
+    # mHC (Manifold-Constrained Hyper-Connections) config
+    mhc_expansion: int = 1
+
+    @property
+    def head_dim(self) -> int:
+        return self.dim // self.n_heads
+
+
+model_config = ModelConfig()
+
+MODEL_SERIES = {
+    "haiku": {
+        "dim": 64,
+        "n_unique_layers": 12,
+        "n_logical_layers": 24,
+        "n_heads": 4,
+        "n_kv_heads": 2,
+        "ffn_dim": 384,
+        "dropout": 0.0,
+        "seq_len": 2048,
+        "mtp_horizons": (2, 3, 4),
+        "rope_fraction": 0.5,
+        "batch_size": 80,
+        "grad_accum": 1,
+        "lr": 8e-4,
+        "min_lr": 1e-5,
+        "sft_lr": 2e-4,
+        "sft_min_lr": 1e-5,
+        "warmup_steps": 300,
+        "weight_decay": 0.02,
+        "pretrain_passes": 2,
+        "sft_passes": 3,
+        "max_sft_target_chars": 0,
+        "use_grad_checkpoint": True,
+        "num_workers": 24,
+        "prefetch_factor": 64,
+        "shuffle_buffer": 8192,
+        "max_pretrain_tokens": 0,
+        "min_pretrain_tokens": 100_000_000,
+        "quantization": "nvfp4",
+        "engram_dim": 8,
+        "engram_heads": 2,
+        "engram_table_size": 64,
+        "engram_max_ngram": 2,
+        "mhc_expansion": 2,
+    },
+    "sonnet": {
+        "dim": 1024,
+        "n_unique_layers": 20,
+        "n_logical_layers": 40,
+        "n_heads": 16,
+        "n_kv_heads": 4,
+        "ffn_dim": 4096,
+        "dropout": 0.0,
+        "seq_len": 2048,
+        "mtp_horizons": (2,),
+        "rope_fraction": 0.5,
+        "batch_size": 24,
+        "grad_accum": 1,
+        "lr": 1e-4,
+        "min_lr": 2e-5,
+        "sft_lr": 5e-5,
+        "sft_min_lr": 5e-6,
+        "warmup_steps": 250,
+        "weight_decay": 0.1,
+        "pretrain_passes": 1,
+        "sft_passes": 1,
+        "max_sft_target_chars": 0,
+        "use_grad_checkpoint": True,
+        "num_workers": 32,
+        "prefetch_factor": 64,
+        "shuffle_buffer": 16384,
+        "max_pretrain_tokens": 0,
+        "min_pretrain_tokens": 100_000_000,
+        "quantization": "nvfp4",
+        "engram_dim": 32,
+        "engram_heads": 8,
+        "engram_table_size": 4096,
+        "engram_max_ngram": 2,
+        "mhc_expansion": 2,
+    },
+    "opus": {
+        "dim": 1536,
+        "n_unique_layers": 18,
+        "n_logical_layers": 36,
+        "n_heads": 16,
+        "n_kv_heads": 4,
+        "ffn_dim": 5888,
+        "dropout": 0.0,
+        "seq_len": 2048,
+        "mtp_horizons": (2,),
+        "rope_fraction": 0.5,
+        "batch_size": 24,
+        "grad_accum": 1,
+        "lr": 1.6e-4,
+        "min_lr": 1.6e-5,
+        "sft_lr": 3e-5,
+        "sft_min_lr": 3e-6,
+        "warmup_steps": 200,
+        "weight_decay": 0.1,
+        "pretrain_passes": 1,
+        "sft_passes": 1,
+        "max_sft_target_chars": 0,
+        "use_grad_checkpoint": True,
+        "num_workers": 48,
+        "prefetch_factor": 64,
+        "shuffle_buffer": 16384,
+        "max_pretrain_tokens": 0,
+        "min_pretrain_tokens": 100_000_000,
+        "quantization": "nvfp4",
+        "engram_dim": 64,
+        "engram_heads": 8,
+        "engram_table_size": 8192,
+        "engram_max_ngram": 2,
+        "mhc_expansion": 4,
+    },
+}
+
+
+# ---------------------------------------------------------------------------
+# Tokenizer (from ailay.tokenizer)
+# ---------------------------------------------------------------------------
+
+FORMAT_TOKENS = [
+    "<|user|>",
+    "<|assistant|>",
+    "<|system|>",
+    "<|start_header_id|>",
+    "<|end_header_id|>",
+    "<|begin_of_thought|>",
+    "<|end_of_thought|>",
+    "<|begin_of_solution|>",
+    "<|end_of_solution|>",
+]
+
+
+class WordTokenizer:
+    WORD_RE = re.compile(
+        r"\s+|[^\W\d_]+(?:['\u2019][^\W\d_]+)?|\d+|[^\w\s]+", re.UNICODE
+    )
+
+    def __init__(
+        self, extra_chars: str = "", format_tokens: Optional[List[str]] = None
+    ) -> None:
+        base = string.ascii_letters + string.digits + string.punctuation + " \n\t\r"
+        fallback_chars = sorted(set(base + extra_chars))
+        self.core_special = ["<PAD>", "<BOS>", "<EOS>", "<UNK>"]
+        self.format_tokens = (
+            list(format_tokens) if format_tokens else list(FORMAT_TOKENS)
+        )
+        self.special = list(self.core_special) + list(self.format_tokens)
+        self.id_to_token: List[str] = (
+            list(self.core_special) + self.format_tokens + fallback_chars
+        )
+        self.token_to_id: Dict[str, int] = {
+            t: i for i, t in enumerate(self.id_to_token)
+        }
+        self.special_multi_tokens = sorted(
+            [t for t in self.special if len(t) > 1], key=len, reverse=True
+        )
+        self.multi_char_tokens = self.special_multi_tokens
+        self.dynamic_additions = 0
+
+    @property
+    def pad_id(self) -> int:
+        return self.token_to_id["<PAD>"]
+
+    @property
+    def bos_id(self) -> int:
+        return self.token_to_id["<BOS>"]
+
+    @property
+    def eos_id(self) -> int:
+        return self.token_to_id["<EOS>"]
+
+    @property
+    def unk_id(self) -> int:
+        return self.token_to_id["<UNK>"]
+
+    @property
+    def vocab_size(self) -> int:
+        return len(self.id_to_token)
+
+    def maybe_add_char(self, ch: str) -> bool:
+        if ch in self.token_to_id:
+            return False
+        self.token_to_id[ch] = len(self.id_to_token)
+        self.id_to_token.append(ch)
+        self.dynamic_additions += 1
+        return True
+
+    def maybe_add_token(self, token: str) -> bool:
+        if token in self.token_to_id:
+            return False
+        self.token_to_id[token] = len(self.id_to_token)
+        self.id_to_token.append(token)
+        self.dynamic_additions += 1
+        return True
+
+    def iter_lexical_tokens(self, text: str) -> Iterator[str]:
+        i = 0
+        n = len(text)
+        while i < n:
+            matched_special = False
+            for token in self.special_multi_tokens:
+                if text.startswith(token, i):
+                    yield token
+                    i += len(token)
+                    matched_special = True
+                    break
+            if matched_special:
+                continue
+            m = self.WORD_RE.match(text, i)
+            if m is None:
+                yield text[i]
+                i += 1
+                continue
+            tok = m.group(0)
+            yield tok
+            i = m.end()
+
+    def encode(
+        self, text: str, add_bos: bool = False, add_eos: bool = False
+    ) -> List[int]:
+        out: List[int] = []
+        if add_bos:
+            out.append(self.bos_id)
+        unk = self.unk_id
+        t2i = self.token_to_id
+        for tok in self.iter_lexical_tokens(text):
+            tid = t2i.get(tok)
+            if tid is not None:
+                out.append(tid)
+                continue
+            for ch in tok:
+                out.append(t2i.get(ch, unk))
+        if add_eos:
+            out.append(self.eos_id)
+        return out
+
+    def decode(self, ids: Sequence[int], skip_special: bool = True) -> str:
+        pieces: List[str] = []
+        for idx in ids:
+            if int(idx) < 0 or int(idx) >= len(self.id_to_token):
+                continue
+            tok = self.id_to_token[int(idx)]
+            if skip_special and tok in self.special:
+                continue
+            pieces.append(tok)
+        return "".join(pieces)
+
+    def save(self, path: Path) -> None:
+        with path.open("w", encoding="utf-8") as f:
+            json.dump(
+                {
+                    "id_to_token": self.id_to_token,
+                    "format_tokens": self.format_tokens,
+                    "core_special": self.core_special,
+                    "tokenizer_type": "word_level_v1",
+                },
+                f,
+                ensure_ascii=False,
+                indent=2,
+            )
+
+    @classmethod
+    def load(cls, path: Path) -> WordTokenizer:
+        with path.open("r", encoding="utf-8") as f:
+            data = json.load(f)
+        format_tokens = data.get("format_tokens", FORMAT_TOKENS)
+        tokenizer = cls(extra_chars="", format_tokens=format_tokens)
+        tokenizer.id_to_token = data["id_to_token"]
+        tokenizer.token_to_id = {t: i for i, t in enumerate(tokenizer.id_to_token)}
+        tokenizer.special = list(tokenizer.core_special) + list(tokenizer.format_tokens)
+        tokenizer.special_multi_tokens = sorted(
+            [t for t in tokenizer.special if len(t) > 1], key=len, reverse=True
+        )
+        tokenizer.multi_char_tokens = tokenizer.special_multi_tokens
+        return tokenizer
+
+
+LetterTokenizer = WordTokenizer
+
+
+# ---------------------------------------------------------------------------
+# Model (from ailay.model)
+# ---------------------------------------------------------------------------
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6) -> None:
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(dim))
+        self.eps = eps
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if hasattr(torch.nn.functional, "rms_norm"):
+            return torch.nn.functional.rms_norm(
+                x, self.weight.shape, self.weight, self.eps
+            )
+        x_fp = x.float()
+        rms = torch.rsqrt(x_fp.pow(2).mean(dim=-1, keepdim=True) + self.eps)
+        return (x_fp * rms).to(dtype=x.dtype) * self.weight
+
+
+class RotaryEmbedding(nn.Module):
+    def __init__(self, dim: int, base: float = 10000.0) -> None:
+        super().__init__()
+        inv = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
+        self.register_buffer("inv_freq", inv, persistent=False)
+
+    def cos_sin(
+        self, seq_len: int, device: torch.device, dtype: torch.dtype
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        t = torch.arange(seq_len, device=device, dtype=self.inv_freq.dtype)
+        freqs = torch.outer(t, self.inv_freq)
+        emb = torch.cat([freqs, freqs], dim=-1)
+        cos = emb.cos()[None, None, :, :].to(dtype=dtype)
+        sin = emb.sin()[None, None, :, :].to(dtype=dtype)
+        return cos, sin
+
+
+def _rotate_half(x: torch.Tensor) -> torch.Tensor:
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+class CausalSelfAttention(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        n_heads: int,
+        n_kv_heads: int,
+        head_dim: int,
+        dropout: float,
+        sliding_window: int,
+        rope_fraction: float,
+    ) -> None:
+        super().__init__()
+        self.dim = dim
+        self.n_heads = n_heads
+        self.n_kv_heads = n_kv_heads
+        self.head_dim = head_dim
+        self.n_rep = n_heads // n_kv_heads
+        self.dropout = dropout
+        self.sliding_window = sliding_window
+
+        self.wq = nn.Linear(dim, n_heads * head_dim, bias=False)
+        self.wk = nn.Linear(dim, n_kv_heads * head_dim, bias=False)
+        self.wv = nn.Linear(dim, n_kv_heads * head_dim, bias=False)
+        self.wo = nn.Linear(n_heads * head_dim, dim, bias=False)
+
+        for lin in (self.wq, self.wk, self.wv):
+            nn.init.normal_(lin.weight, std=dim ** -0.5)
+        nn.init.normal_(self.wo.weight, std=(n_heads * head_dim) ** -0.5)
+
+        self.rope_dim = max(2, int(head_dim * rope_fraction) // 2 * 2)
+        self.rope = RotaryEmbedding(self.rope_dim)
+
+        self.q_norm = RMSNorm(head_dim)
+        self.k_norm = RMSNorm(head_dim)
+
+        self.output_gate = nn.Parameter(torch.zeros(n_heads))
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        is_global: bool,
+        past_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        use_cache: bool = False,
+    ) -> Tuple[torch.Tensor, Optional[Tuple[torch.Tensor, torch.Tensor]]]:
+        B, T, _ = x.shape
+
+        q = self.wq(x).view(B, T, self.n_heads, self.head_dim)
+        k = self.wk(x).view(B, T, self.n_kv_heads, self.head_dim)
+        v = self.wv(x).view(B, T, self.n_kv_heads, self.head_dim)
+
+        q = self.q_norm(q)
+        k = self.k_norm(k)
+
+        q = q.transpose(1, 2)
+        k = k.transpose(1, 2)
+        v = v.transpose(1, 2)
+
+        past_len = past_kv[0].shape[2] if past_kv is not None else 0
+        cos, sin = self.rope.cos_sin(T + past_len, x.device, q.dtype)
+        cos_slice = cos[:, :, past_len : past_len + T, :]
+        sin_slice = sin[:, :, past_len : past_len + T, :]
+
+        q_rope = q[..., : self.rope_dim]
+        q_pass = q[..., self.rope_dim :]
+        k_rope = k[..., : self.rope_dim]
+        k_pass = k[..., self.rope_dim :]
+
+        q_rope = (q_rope * cos_slice) + (_rotate_half(q_rope) * sin_slice)
+        k_rope = (k_rope * cos_slice) + (_rotate_half(k_rope) * sin_slice)
+
+        q = torch.cat([q_rope, q_pass], dim=-1)
+        k = torch.cat([k_rope, k_pass], dim=-1)
+
+        if past_kv is not None:
+            k = torch.cat([past_kv[0], k], dim=2)
+            v = torch.cat([past_kv[1], v], dim=2)
+
+        new_kv = (k, v) if use_cache else None
+
+        S = k.shape[2]
+        if self.n_rep > 1:
+            k = (
+                k[:, :, None, :, :]
+                .expand(B, self.n_kv_heads, self.n_rep, S, self.head_dim)
+                .reshape(B, self.n_heads, S, self.head_dim)
+            )
+            v = (
+                v[:, :, None, :, :]
+                .expand(B, self.n_kv_heads, self.n_rep, S, self.head_dim)
+                .reshape(B, self.n_heads, S, self.head_dim)
+            )
+
+        drop_p = self.dropout if (self.training and torch.is_grad_enabled()) else 0.0
+
+        if is_global:
+            if past_kv is None and T > 1:
+                out = F.scaled_dot_product_attention(
+                    q, k, v, is_causal=True, dropout_p=drop_p
+                )
+            else:
+                out = F.scaled_dot_product_attention(q, k, v, dropout_p=drop_p)
+        else:
+            T_q = q.shape[2]
+            q_pos = torch.arange(past_len, past_len + T_q, device=q.device).unsqueeze(1)
+            k_pos = torch.arange(S, device=q.device).unsqueeze(0)
+            mask = (q_pos >= k_pos) & ((q_pos - k_pos) < self.sliding_window)
+            out = F.scaled_dot_product_attention(
+                q, k, v, attn_mask=mask.unsqueeze(0).unsqueeze(0), dropout_p=drop_p
+            )
+
+        gate = torch.sigmoid(self.output_gate).view(1, self.n_heads, 1, 1)
+        out = out * gate
+
+        out = out.transpose(1, 2).contiguous().view(B, T, self.n_heads * self.head_dim)
+        out = self.wo(out)
+
+        return out, new_kv
+
+
+class SwiGLU(nn.Module):
+    def __init__(self, dim: int, hidden_dim: int, dropout: float) -> None:
+        super().__init__()
+        self.gate = nn.Linear(dim, hidden_dim, bias=False)
+        self.up = nn.Linear(dim, hidden_dim, bias=False)
+        self.down = nn.Linear(hidden_dim, dim, bias=False)
+        self.drop = nn.Dropout(dropout)
+
+        nn.init.normal_(self.gate.weight, std=dim ** -0.5)
+        nn.init.normal_(self.up.weight, std=dim ** -0.5)
+        nn.init.normal_(self.down.weight, std=hidden_dim ** -0.5)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        h = F.silu(self.gate(x)) * self.up(x)
+        out = self.down(h)
+        if self.training and torch.is_grad_enabled():
+            out = self.drop(out)
+        return out
+
+
+class EngramBlock(nn.Module):
+    """Conditional memory via O(1) hashed N-gram lookup (DeepSeek Engram)."""
+
+    def __init__(
+        self,
+        dim: int,
+        engram_dim: int,
+        n_heads: int = 4,
+        table_size: int = 8192,
+        max_ngram: int = 3,
+    ) -> None:
+        super().__init__()
+        self.dim = dim
+        self.engram_dim = engram_dim
+        self.n_heads = n_heads
+        self.table_size = table_size
+        self.max_ngram = max_ngram
+
+        self.embeddings = nn.ParameterDict()
+        for n in range(2, max_ngram + 1):
+            for k in range(n_heads):
+                self.embeddings[f"{n}_{k}"] = nn.Parameter(
+                    torch.randn(table_size, engram_dim) * (engram_dim**-0.5)
+                )
+
+        for n in range(2, max_ngram + 1):
+            for k in range(n_heads):
+                seed = int(hashlib.md5(f"engram_{n}_{k}".encode()).hexdigest()[:8], 16)
+                rng = torch.Generator().manual_seed(seed)
+                a = torch.randint(1, 2**31, (1,), generator=rng).item()
+                b = torch.randint(0, 2**31, (1,), generator=rng).item()
+                self.register_buffer(
+                    f"hash_a_{n}_{k}", torch.tensor(a), persistent=False
+                )
+                self.register_buffer(
+                    f"hash_b_{n}_{k}", torch.tensor(b), persistent=False
+                )
+
+        total_branch_dim = engram_dim * n_heads * (max_ngram - 1)
+        self.branch_conv = nn.Conv1d(
+            total_branch_dim,
+            total_branch_dim,
+            kernel_size=4,
+            dilation=max_ngram,
+            padding=0,
+            groups=total_branch_dim,
+            bias=True,
+        )
+        nn.init.zeros_(self.branch_conv.weight)
+        nn.init.zeros_(self.branch_conv.bias)
+
+        self.gate_query = nn.Linear(dim, engram_dim, bias=False)
+        self.gate_key = nn.Linear(total_branch_dim, engram_dim, bias=False)
+        self.gate_value = nn.Linear(total_branch_dim, dim, bias=False)
+        self.gate_scale = engram_dim**-0.5
+
+    def _hash_ngram(self, token_ids: torch.Tensor, n: int, k: int) -> torch.Tensor:
+        a = getattr(self, f"hash_a_{n}_{k}")
+        b = getattr(self, f"hash_b_{n}_{k}")
+        B, T = token_ids.shape
+        padded = F.pad(token_ids, (n - 1, 0), value=0)
+        combined = torch.zeros(B, T, dtype=torch.long, device=token_ids.device)
+        for i in range(n):
+            combined = (combined * 31 + padded[:, i : i + T].long()) % self.table_size
+        return ((a * combined) ^ b) % self.table_size
+
+    def forward(
+        self, hidden: torch.Tensor, token_ids: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        B, T, _ = hidden.shape
+        if token_ids is None:
+            token_ids = hidden.mean(dim=-1).long() % self.table_size
+        all_indices = []
+        all_tables = []
+        for n in range(2, self.max_ngram + 1):
+            for k in range(self.n_heads):
+                all_indices.append(self._hash_ngram(token_ids, n, k))
+                all_tables.append(self.embeddings[f"{n}_{k}"])
+        branch_outputs = [tbl[idx] for idx, tbl in zip(all_indices, all_tables)]
+        memory = torch.cat(branch_outputs, dim=-1)
+        conv_in = memory.transpose(1, 2)
+        conv_in = F.pad(
+            conv_in,
+            (self.branch_conv.dilation[0] * (self.branch_conv.kernel_size[0] - 1), 0),
+        )
+        conv_out = self.branch_conv(conv_in)
+        memory = conv_out.transpose(1, 2)
+        query = self.gate_query(hidden)
+        key = self.gate_key(memory)
+        gate = torch.sigmoid(
+            (query * key).sum(dim=-1, keepdim=True) * self.gate_scale
+        )
+        value = self.gate_value(memory)
+        return gate * value
+
+
+def _sinkhorn_knopp(logits: torch.Tensor, n_iters: int = 7) -> torch.Tensor:
+    M = torch.exp(logits.clamp(-10, 10))
+    for _ in range(n_iters):
+        M = M / M.sum(dim=-1, keepdim=True).clamp(min=1e-10)
+        M = M / M.sum(dim=-2, keepdim=True).clamp(min=1e-10)
+    return M
+
+
+class ManifoldHyperConnection(nn.Module):
+    """Manifold-Constrained Hyper-Connections (mHC) residual wrapper."""
+
+    def __init__(self, dim: int, expansion: int = 2) -> None:
+        super().__init__()
+        self.dim = dim
+        self.expansion = expansion
+        n = expansion
+
+        self.bias_pre = nn.Parameter(torch.zeros(1, n))
+        self.bias_post = nn.Parameter(torch.zeros(1, n))
+        self.bias_res = nn.Parameter(torch.zeros(n, n))
+
+        self.theta_pre = nn.Linear(n * dim, n, bias=False)
+        self.theta_post = nn.Linear(n * dim, n, bias=False)
+        self.theta_res = nn.Linear(n * dim, n * n, bias=False)
+
+        self.alpha_pre = nn.Parameter(torch.tensor(0.0))
+        self.alpha_post = nn.Parameter(torch.tensor(0.0))
+        self.alpha_res = nn.Parameter(torch.tensor(0.0))
+
+    def _compute_mappings(
+        self, x_expanded: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        B, T, _ = x_expanded.shape
+        n = self.expansion
+        x_norm = F.rms_norm(x_expanded, [x_expanded.shape[-1]])
+        d_pre = torch.tanh(self.theta_pre(x_norm))
+        d_post = torch.tanh(self.theta_post(x_norm))
+        d_res = self.theta_res(x_norm)
+        H_pre_raw = torch.sigmoid(self.alpha_pre * d_pre + self.bias_pre)
+        H_post_raw = 2.0 * torch.sigmoid(self.alpha_post * d_post + self.bias_post)
+        H_res_raw = (self.alpha_res * d_res + self.bias_res.reshape(1, 1, -1)).reshape(
+            B, T, n, n
+        )
+        H_res = _sinkhorn_knopp(H_res_raw)
+        return H_pre_raw.unsqueeze(-2), H_post_raw.unsqueeze(-2), H_res
+
+    def expand_stream(self, x: torch.Tensor) -> torch.Tensor:
+        return x.repeat(1, 1, self.expansion)
+
+    def collapse_stream(self, x_expanded: torch.Tensor) -> torch.Tensor:
+        B, T, _ = x_expanded.shape
+        return x_expanded.view(B, T, self.expansion, self.dim).mean(dim=-2)
+
+    def pre_mix(self, x_expanded: torch.Tensor, H_pre: torch.Tensor) -> torch.Tensor:
+        B, T, _ = x_expanded.shape
+        x_streams = x_expanded.view(B, T, self.expansion, self.dim)
+        return (H_pre @ x_streams).squeeze(-2)
+
+    def post_res_mix(
+        self,
+        layer_output: torch.Tensor,
+        x_expanded: torch.Tensor,
+        H_post: torch.Tensor,
+        H_res: torch.Tensor,
+    ) -> torch.Tensor:
+        B, T, _ = x_expanded.shape
+        x_streams = x_expanded.view(B, T, self.expansion, self.dim)
+        mixed = torch.matmul(H_res, x_streams)
+        post_out = torch.matmul(H_post.transpose(-2, -1), layer_output.unsqueeze(-2))
+        return (mixed + post_out).reshape(B, T, self.expansion * self.dim)
+
+
+class TransformerBlock(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        n_heads: int,
+        n_kv_heads: int,
+        head_dim: int,
+        ffn_dim: int,
+        dropout: float,
+        sliding_window: int,
+        rope_fraction: float,
+        engram_dim: int = 0,
+        engram_heads: int = 4,
+        engram_table_size: int = 8192,
+        engram_max_ngram: int = 3,
+        mhc_expansion: int = 1,
+    ) -> None:
+        super().__init__()
+        self.dim = dim
+        self.norm1 = RMSNorm(dim)
+        self.attn = CausalSelfAttention(
+            dim=dim,
+            n_heads=n_heads,
+            n_kv_heads=n_kv_heads,
+            head_dim=head_dim,
+            dropout=dropout,
+            sliding_window=sliding_window,
+            rope_fraction=rope_fraction,
+        )
+        self.norm2 = RMSNorm(dim)
+        self.ffn = SwiGLU(dim, ffn_dim, dropout)
+        self.use_engram = engram_dim > 0
+        if self.use_engram:
+            self.engram = EngramBlock(
+                dim=dim,
+                engram_dim=engram_dim,
+                n_heads=engram_heads,
+                table_size=engram_table_size,
+                max_ngram=engram_max_ngram,
+            )
+            self.engram_norm = RMSNorm(dim)
+        self.use_mhc = mhc_expansion > 1
+        if self.use_mhc:
+            self.mhc_attn = ManifoldHyperConnection(dim, expansion=mhc_expansion)
+            self.mhc_ffn = ManifoldHyperConnection(dim, expansion=mhc_expansion)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        is_global: bool,
+        past_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        use_cache: bool = False,
+        token_ids: Optional[torch.Tensor] = None,
+    ) -> Tuple[torch.Tensor, Optional[Tuple[torch.Tensor, torch.Tensor]]]:
+        if self.use_mhc:
+            x_exp = self.mhc_attn.expand_stream(x)
+            H_pre, H_post, H_res = self.mhc_attn._compute_mappings(x_exp)
+            attn_in = self.mhc_attn.pre_mix(x_exp, H_pre)
+            attn_out, new_kv = self.attn(
+                self.norm1(attn_in), is_global, past_kv, use_cache
+            )
+            x_exp = self.mhc_attn.post_res_mix(attn_out, x_exp, H_post, H_res)
+            if self.use_engram:
+                collapsed = self.mhc_attn.collapse_stream(x_exp)
+                collapsed = collapsed + self.engram(
+                    self.engram_norm(collapsed), token_ids=token_ids
+                )
+                x_exp = self.mhc_attn.expand_stream(collapsed)
+            H_pre2, H_post2, H_res2 = self.mhc_ffn._compute_mappings(x_exp)
+            ffn_in = self.mhc_ffn.pre_mix(x_exp, H_pre2)
+            ffn_out = self.ffn(self.norm2(ffn_in))
+            x_exp = self.mhc_ffn.post_res_mix(ffn_out, x_exp, H_post2, H_res2)
+            x = self.mhc_attn.collapse_stream(x_exp)
+        else:
+            attn_out, new_kv = self.attn(self.norm1(x), is_global, past_kv, use_cache)
+            x = x + attn_out
+            if self.use_engram:
+                x = x + self.engram(self.engram_norm(x), token_ids=token_ids)
+            x = x + self.ffn(self.norm2(x))
+        return x, new_kv
+
+
+def _detect_engram_dim(state_dict: dict) -> int:
+    for key in state_dict:
+        if ".engram." in key and ".embeddings." in key:
+            return state_dict[key].shape[-1]
+    return 0
+
+
+def _detect_mhc_expansion(state_dict: dict) -> int:
+    for key, val in state_dict.items():
+        if ".mhc_attn.bias_pre" in key and val.dim() == 2:
+            return val.shape[-1]
+    return 1
+
+
+class TinyMemoryLM(nn.Module):
+    def __init__(
+        self,
+        vocab_size: int,
+        dim: int,
+        n_unique_layers: int,
+        n_logical_layers: int,
+        n_heads: int,
+        n_kv_heads: int,
+        ffn_dim: int,
+        dropout: float,
+        mtp_horizons: Sequence[int],
+        grad_checkpoint: bool,
+        sliding_window: int = 512,
+        rope_fraction: float = 0.5,
+        embed_scale: bool = True,
+        engram_dim: int = 0,
+        engram_heads: int = 4,
+        engram_table_size: int = 8192,
+        engram_max_ngram: int = 3,
+        mhc_expansion: int = 1,
+    ) -> None:
+        super().__init__()
+        self.dim = dim
+        self.n_unique_layers = n_unique_layers
+        self.n_logical_layers = n_logical_layers
+        self.grad_checkpoint = grad_checkpoint
+        self.embed_scale_factor = math.sqrt(dim) if embed_scale else 1.0
+        head_dim = dim // n_heads
+
+        self.embed_tokens = nn.Embedding(vocab_size, dim)
+        self.head = nn.Linear(dim, vocab_size, bias=False)
+        self.head.weight = self.embed_tokens.weight
+
+        self.output_bias = nn.Parameter(torch.zeros(vocab_size))
+
+        self.blocks = nn.ModuleList(
+            [
+                TransformerBlock(
+                    dim=dim,
+                    n_heads=n_heads,
+                    n_kv_heads=n_kv_heads,
+                    head_dim=head_dim,
+                    ffn_dim=ffn_dim,
+                    dropout=dropout,
+                    sliding_window=sliding_window,
+                    rope_fraction=rope_fraction,
+                    engram_dim=engram_dim,
+                    engram_heads=engram_heads,
+                    engram_table_size=engram_table_size,
+                    engram_max_ngram=engram_max_ngram,
+                    mhc_expansion=mhc_expansion,
+                )
+                for _ in range(n_unique_layers)
+            ]
+        )
+        self.norm = RMSNorm(dim)
+
+        self.mtp_horizons = sorted({int(h) for h in mtp_horizons if int(h) > 1})
+        self.mtp_adapters = nn.ModuleDict(
+            {str(h): nn.Linear(dim, dim, bias=False) for h in self.mtp_horizons}
+        )
+        self.mtp_norms = nn.ModuleDict(
+            {str(h): RMSNorm(dim) for h in self.mtp_horizons}
+        )
+
+        res_scale = (2 * n_logical_layers) ** -0.5
+        for block in self.blocks:
+            block.attn.wo.weight.data.mul_(res_scale)
+            block.ffn.down.weight.data.mul_(res_scale)
+
+    def resize_token_embeddings(self, new_vocab_size: int) -> None:
+        old_vocab_size = self.embed_tokens.num_embeddings
+        if new_vocab_size == old_vocab_size:
+            return
+        device = self.embed_tokens.weight.device
+        old_embed_weight = self.embed_tokens.weight.data.clone()
+        self.embed_tokens = nn.Embedding(
+            new_vocab_size, self.embed_tokens.embedding_dim
+        ).to(device)
+        self.head = nn.Linear(
+            self.embed_tokens.embedding_dim, new_vocab_size, bias=False
+        ).to(device)
+        self.head.weight = self.embed_tokens.weight
+        old_bias = self.output_bias.data.clone()
+        self.output_bias = nn.Parameter(torch.zeros(new_vocab_size, device=device))
+        copy_size = min(old_vocab_size, new_vocab_size)
+        self.output_bias.data[:copy_size] = old_bias[:copy_size]
+        self.embed_tokens.weight.data[:copy_size] = old_embed_weight[:copy_size]
+
+    def _build_logical_layers(self) -> List[Tuple[nn.Module, int]]:
+        logical = []
+        blocks_list = list(self.blocks)
+        full_sequence = blocks_list + blocks_list
+        for logical_idx, block in enumerate(full_sequence[: self.n_logical_layers]):
+            logical.append((block, logical_idx))
+        return logical
+
+    def forward(
+        self,
+        ids: torch.Tensor,
+        use_cache: bool = False,
+        past_key_values: Optional[
+            List[Optional[Tuple[torch.Tensor, torch.Tensor]]]
+        ] = None,
+        return_hidden: bool = False,
+    ) -> Tuple[
+        torch.Tensor,
+        Dict[int, torch.Tensor],
+        Optional[torch.Tensor],
+        Optional[List[Tuple[torch.Tensor, torch.Tensor]]],
+    ]:
+        B, T = ids.shape
+        x = self.embed_tokens(ids) * self.embed_scale_factor
+
+        logical_layers = self._build_logical_layers()
+        new_past_key_values: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = (
+            [] if use_cache else None
+        )
+
+        for layer_idx, (block, logical_idx) in enumerate(logical_layers):
+            is_global = logical_idx % 2 == 0
+            past_kv = (
+                past_key_values[layer_idx]
+                if past_key_values is not None and layer_idx < len(past_key_values)
+                else None
+            )
+
+            if self.grad_checkpoint and self.training and not use_cache:
+                x, layer_kv = checkpoint(
+                    block, x, is_global, past_kv, use_cache, ids, use_reentrant=True
+                )
+            else:
+                x, layer_kv = block(x, is_global, past_kv, use_cache, ids)
+
+            if new_past_key_values is not None:
+                new_past_key_values.append(layer_kv)
+
+        x = self.norm(x)
+        h_out = x if return_hidden else None
+        logits = self.head(x)
+        if self.embed_scale_factor != 1.0:
+            logits = logits / self.embed_scale_factor
+        logits = logits + self.output_bias
+
+        mtp: Dict[int, torch.Tensor] = {}
+        if self.mtp_horizons and self.training:
+            for horizon in self.mtp_horizons:
+                if horizon > 1 and horizon <= T - 1:
+                    shifted_h = x[:, :-horizon, :]
+                    adapted_h = self.mtp_adapters[str(horizon)](shifted_h)
+                    adapted_h = self.mtp_norms[str(horizon)](adapted_h)
+                    mtp_logits = self.head(adapted_h)
+                    if self.embed_scale_factor != 1.0:
+                        mtp_logits = mtp_logits / self.embed_scale_factor
+                    mtp_logits = mtp_logits + self.output_bias
+                    mtp[horizon] = mtp_logits
+
+        return logits, mtp, h_out, new_past_key_values
+
+
+# ---------------------------------------------------------------------------
+# Generation (from ailay.generation)
+# ---------------------------------------------------------------------------
+
+
+def sample_text(
+    model: TinyMemoryLM,
+    tokenizer: WordTokenizer,
+    prompt: str,
+    max_new_tokens: int,
+    temperature: float,
+    top_k: int,
+    branches: int,
+    branch_len: int,
+    device: torch.device,
+    seq_len: int,
+) -> str:
+    def _sample_id(logits: torch.Tensor) -> torch.Tensor:
+        if not torch.isfinite(logits).any():
+            logits = torch.zeros_like(logits)
+        logits = torch.where(
+            torch.isfinite(logits), logits, torch.full_like(logits, -1e9)
+        )
+        if top_k > 0:
+            v, idx = torch.topk(logits, k=min(top_k, logits.shape[-1]))
+            p = torch.softmax(v, dim=-1)
+            return idx.gather(-1, torch.multinomial(p, 1))
+        p = torch.softmax(logits, dim=-1)
+        return torch.multinomial(p, 1)
+
+    model.eval()
+    ids = tokenizer.encode(prompt, add_bos=True, add_eos=False)
+    prompt_len = len(ids)
+    x = torch.tensor(ids, dtype=torch.long, device=device).unsqueeze(0)
+
+    with torch.no_grad():
+        generated = 0
+        while generated < max_new_tokens:
+            if branches <= 1:
+                ctx = x[:, -seq_len:]
+                logits, _, _, _ = model(ctx)
+                nlogits = logits[:, -1, :] / max(temperature, 1e-6)
+                nid = _sample_id(nlogits)
+                x = torch.cat([x, nid], dim=1)
+                generated += 1
+                continue
+            rollout = min(branch_len, max_new_tokens - generated)
+            best_nll: Optional[float] = None
+            best_tokens: Optional[List[torch.Tensor]] = None
+            for _ in range(branches):
+                cand = x
+                cand_tokens: List[torch.Tensor] = []
+                nll = 0.0
+                for _ in range(rollout):
+                    ctx = cand[:, -seq_len:]
+                    logits, _, _, _ = model(ctx)
+                    nlogits = logits[:, -1, :] / max(temperature, 1e-6)
+                    nid = _sample_id(nlogits)
+                    lp = F.log_softmax(nlogits.float(), dim=-1)
+                    nll += float(-lp.gather(-1, nid).item())
+                    cand = torch.cat([cand, nid], dim=1)
+                    cand_tokens.append(nid)
+                if best_nll is None or nll < best_nll:
+                    best_nll = nll
+                    best_tokens = cand_tokens
+            assert best_tokens is not None
+            for t in best_tokens:
+                x = torch.cat([x, t], dim=1)
+                generated += 1
+
+    generated_ids = x[0, prompt_len:].tolist()
+    return tokenizer.decode(generated_ids, skip_special=True)
+
+
+def sample_text_cached(
+    model: TinyMemoryLM,
+    tokenizer: WordTokenizer,
+    prompt: str,
+    max_new_tokens: int,
+    temperature: float,
+    top_k: int,
+    device: torch.device,
+    seq_len: int,
+) -> str:
+    model.eval()
+    ids = tokenizer.encode(prompt, add_bos=True, add_eos=False)
+    prompt_len = len(ids)
+    x = torch.tensor(ids, dtype=torch.long, device=device).unsqueeze(0)
+
+    with torch.no_grad():
+        logits, _, _, past_kv = model(x, use_cache=True)
+        nlogits = logits[:, -1, :] / max(temperature, 1e-6)
+        if top_k > 0:
+            v, idx = torch.topk(nlogits, k=min(top_k, nlogits.shape[-1]))
+            p = torch.softmax(v, dim=-1)
+            nid = idx.gather(-1, torch.multinomial(p, 1))
+        else:
+            p = torch.softmax(nlogits, dim=-1)
+            nid = torch.multinomial(p, 1)
+        all_ids = [int(nid.item())]
+
+        for _ in range(max_new_tokens - 1):
+            logits, _, _, past_kv = model(nid, use_cache=True, past_key_values=past_kv)
+            nlogits = logits[:, -1, :] / max(temperature, 1e-6)
+            if top_k > 0:
+                v, idx = torch.topk(nlogits, k=min(top_k, nlogits.shape[-1]))
+                p = torch.softmax(v, dim=-1)
+                nid = idx.gather(-1, torch.multinomial(p, 1))
+            else:
+                p = torch.softmax(nlogits, dim=-1)
+                nid = torch.multinomial(p, 1)
+            tid = int(nid.item())
+            all_ids.append(tid)
+            if tid == tokenizer.eos_id:
+                break
+
+    return tokenizer.decode(all_ids, skip_special=True)
+
+
+def speculative_decode(
+    model: TinyMemoryLM,
+    tokenizer: WordTokenizer,
+    prompt: str,
+    max_new_tokens: int,
+    temperature: float,
+    top_k: int,
+    device: torch.device,
+    seq_len: int,
+) -> str:
+    model.eval()
+    ids = tokenizer.encode(prompt, add_bos=True, add_eos=False)
+    x = torch.tensor(ids, dtype=torch.long, device=device).unsqueeze(0)
+    all_generated: List[int] = []
+
+    with torch.no_grad():
+        logits, _, h_out, past_kv = model(x, use_cache=True, return_hidden=True)
+
+        def _sample_from(lg: torch.Tensor) -> int:
+            lg = lg / max(temperature, 1e-6)
+            if top_k > 0:
+                v, idx = torch.topk(lg, k=min(top_k, lg.shape[-1]))
+                p = torch.softmax(v, dim=-1)
+                return int(idx[torch.multinomial(p, 1)].item())
+            p = torch.softmax(lg, dim=-1)
+            return int(torch.multinomial(p, 1).item())
+
+        main_token = _sample_from(logits[0, -1, :])
+        all_generated.append(main_token)
+
+        while len(all_generated) < max_new_tokens:
+            if main_token == tokenizer.eos_id:
+                break
+
+            draft_tokens = []
+            if h_out is not None and model.mtp_horizons:
+                last_hidden = h_out[:, -1:, :]
+                for h in model.mtp_horizons:
+                    adapter = model.mtp_adapters[str(h)]
+                    norm = model.mtp_norms[str(h)]
+                    adapted = norm(adapter(last_hidden))
+                    draft_logits = model.head(adapted) + model.output_bias
+                    draft_tok = _sample_from(draft_logits[0, 0, :])
+                    draft_tokens.append(draft_tok)
+
+            if not draft_tokens:
+                nid = torch.tensor([[main_token]], dtype=torch.long, device=device)
+                logits, _, h_out, past_kv = model(
+                    nid, use_cache=True, past_key_values=past_kv, return_hidden=True
+                )
+                main_token = _sample_from(logits[0, -1, :])
+                all_generated.append(main_token)
+                continue
+
+            verify_input = torch.tensor(
+                [[main_token] + draft_tokens], dtype=torch.long, device=device
+            )
+            verify_logits, _, h_out, past_kv = model(
+                verify_input,
+                use_cache=True,
+                past_key_values=past_kv,
+                return_hidden=True,
+            )
+
+            accepted = 0
+            all_generated.append(main_token) if main_token not in all_generated[
+                -1:
+            ] else None
+            for i, draft_tok in enumerate(draft_tokens):
+                verified_tok = _sample_from(verify_logits[0, i, :])
+                if verified_tok == draft_tok:
+                    all_generated.append(draft_tok)
+                    accepted += 1
+                    if draft_tok == tokenizer.eos_id:
+                        break
+                else:
+                    all_generated.append(verified_tok)
+                    break
+
+            if accepted < len(draft_tokens):
+                trim_len = len(draft_tokens) - accepted - 1
+                if trim_len > 0 and past_kv is not None:
+                    past_kv = [
+                        (k[:, :, :-trim_len, :], v[:, :, :-trim_len, :])
+                        if k is not None
+                        else None
+                        for k, v in past_kv
+                    ]
+
+            main_token = all_generated[-1]
+
+    return tokenizer.decode(all_generated, skip_special=True)
+
+
+def build_stop_token_ids(tokenizer: WordTokenizer) -> set:
+    stop_tokens = {tokenizer.eos_id}
+    for tok in ("<|user|>", "<|system|>", "<|assistant|>"):
+        tid = tokenizer.token_to_id.get(tok)
+        if tid is not None:
+            stop_tokens.add(int(tid))
+    return stop_tokens
+
+
+def apply_no_repeat_ngram(
+    logits: torch.Tensor,
+    token_history: Sequence[int],
+    ngram_size: int,
+) -> torch.Tensor:
+    if ngram_size <= 1 or len(token_history) < max(0, ngram_size - 1):
+        return logits
+    prefix = tuple(token_history[-(ngram_size - 1) :]) if ngram_size > 1 else tuple()
+    banned: set = set()
+    for i in range(len(token_history) - ngram_size + 1):
+        if tuple(token_history[i : i + ngram_size - 1]) == prefix:
+            banned.add(int(token_history[i + ngram_size - 1]))
+    if not banned:
+        return logits
+    out = logits.clone()
+    banned_ids = torch.tensor(sorted(banned), device=logits.device, dtype=torch.long)
+    out[banned_ids] = float("-inf")
+    return out
+
+
+def score_candidate(
+    prompt: str,
+    raw_text: str,
+    visible_text: str,
+    avg_logprob: float,
+) -> float:
+    clean = visible_text.strip()
+    if not clean:
+        return -1e9
+    score = avg_logprob
+    words = clean.lower().split()
+    prompt_words = re.findall(r"[A-Za-z][A-Za-z'-]{2,}", prompt.lower())
+    prompt_stop = {
+        "what",
+        "which",
+        "when",
+        "where",
+        "why",
+        "how",
+        "are",
+        "is",
+        "the",
+        "and",
+        "for",
+        "with",
+        "that",
+        "this",
+        "from",
+        "into",
+        "about",
+        "explain",
+        "tell",
+        "give",
+        "list",
+        "show",
+        "write",
+        "their",
+        "there",
+        "your",
+    }
+    prompt_keywords = {w for w in prompt_words if w not in prompt_stop}
+    candidate_keywords = set(re.findall(r"[A-Za-z][A-Za-z'-]{2,}", clean.lower()))
+    if len(words) < 6:
+        score -= 2.0
+    else:
+        score += min(2.0, len(words) * 0.03)
+    if clean[-1:] in ".!?":
+        score += 0.5
+    if "<|user|>" in raw_text or "<|system|>" in raw_text:
+        score -= 4.0
+    if raw_text.count("<|assistant|>") > 1:
+        score -= 2.0
+    if prompt_keywords:
+        overlap = len(prompt_keywords & candidate_keywords) / len(prompt_keywords)
+        if overlap == 0.0:
+            score -= 2.5
+        else:
+            score += min(3.5, overlap * 4.0)
+    for open_tok, close_tok in [
+        ("<|begin_of_thought|>", "<|end_of_thought|>"),
+        ("<|begin_of_solution|>", "<|end_of_solution|>"),
+    ]:
+        if (open_tok in raw_text) != (close_tok in raw_text):
+            score -= 1.0
+    if len(words) >= 3:
+        trigrams = [tuple(words[i : i + 3]) for i in range(len(words) - 2)]
+        if trigrams:
+            unique_ratio = len(set(trigrams)) / len(trigrams)
+            if unique_ratio < 0.35:
+                score -= 4.0
+            elif unique_ratio < 0.55:
+                score -= 2.0
+            else:
+                score += min(1.0, (unique_ratio - 0.55) * 2.0)
+    alpha_words = [
+        w
+        for w in words
+        if len(w) <= 18 and (sum(ch.isalpha() for ch in w) / max(len(w), 1)) > 0.7
+    ]
+    alpha_ratio = len(alpha_words) / max(len(words), 1)
+    if alpha_ratio < 0.45:
+        score -= 3.0
+    elif alpha_ratio < 0.65:
+        score -= 1.0
+    return score
+
+
+def generate_candidate(
+    model: TinyMemoryLM,
+    tokenizer: WordTokenizer,
+    prompt: str,
+    max_new_tokens: int,
+    temperature: float,
+    top_k: int,
+    repetition_penalty: float,
+    no_repeat_ngram_size: int,
+    device: str,
+    sft_mode: bool,
+    force_thought: bool,
+    stream: bool,
+    context_window: int,
+) -> Tuple[str, str, float, int]:
+    if sft_mode:
+        full_prompt = f"<|user|>\n{prompt}\n<|assistant|>\n"
+    else:
+        full_prompt = prompt
+    if force_thought:
+        full_prompt = f"{full_prompt}<|begin_of_thought|> "
+    input_ids = tokenizer.encode(full_prompt, add_bos=True, add_eos=False)
+    input_ids_t = torch.tensor([input_ids], dtype=torch.long, device=device)
+    visible_tokens: List[str] = []
+    raw_tokens: List[str] = []
+    stop_token_ids = build_stop_token_ids(tokenizer)
+    total_logprob = 0.0
+    sampled_tokens = 0
+    with torch.no_grad():
+        for _ in range(max_new_tokens):
+            ctx_ids = (
+                input_ids_t[:, -context_window:] if context_window > 0 else input_ids_t
+            )
+            logits, _, _, _ = model(ctx_ids)
+            next_logits = logits[0, -1, :].clone()
+            raw_next_logits = next_logits.clone()
+            if repetition_penalty != 1.0:
+                seen = set(input_ids_t[0].tolist())
+                for token_id in seen:
+                    if next_logits[token_id] > 0:
+                        next_logits[token_id] /= repetition_penalty
+                    else:
+                        next_logits[token_id] *= repetition_penalty
+            if temperature != 1.0:
+                next_logits = next_logits / max(temperature, 1e-6)
+            if no_repeat_ngram_size > 1:
+                next_logits = apply_no_repeat_ngram(
+                    next_logits,
+                    input_ids_t[0].tolist(),
+                    no_repeat_ngram_size,
+                )
+            if top_k > 0:
+                v, _ = torch.topk(next_logits, min(top_k, next_logits.size(0)))
+                next_logits[next_logits < v[-1]] = float("-inf")
+            top_p = 0.9
+            if top_p < 1.0:
+                sorted_logits, sorted_indices = torch.sort(next_logits, descending=True)
+                cum_probs = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1)
+                remove_mask = cum_probs - torch.softmax(sorted_logits, dim=-1) >= top_p
+                sorted_logits[remove_mask] = float("-inf")
+                next_logits = sorted_logits.scatter(0, sorted_indices, sorted_logits)
+            if not torch.isfinite(next_logits).any():
+                next_logits = raw_next_logits
+                if temperature != 1.0:
+                    next_logits = next_logits / max(temperature, 1e-6)
+            probs = torch.softmax(next_logits, dim=-1)
+            next_id = torch.multinomial(probs, num_samples=1).item()
+            total_logprob += float(torch.log(probs[next_id] + 1e-12).item())
+            sampled_tokens += 1
+            if next_id in stop_token_ids:
+                break
+            token_str = (
+                tokenizer.id_to_token[next_id]
+                if next_id < len(tokenizer.id_to_token)
+                else ""
+            )
+            raw_tokens.append(token_str)
+            if token_str not in tokenizer.special:
+                visible_tokens.append(token_str)
+                if stream:
+                    print(token_str, end="", flush=True)
+            input_ids_t = torch.cat(
+                [input_ids_t, torch.tensor([[next_id]], device=device)], dim=1
+            )
+    if stream:
+        print()
+    avg_logprob = total_logprob / max(1, sampled_tokens)
+    return "".join(visible_tokens), "".join(raw_tokens), avg_logprob, 0
+
+
+def generate_beam_search(
+    model: TinyMemoryLM,
+    tokenizer: WordTokenizer,
+    prompt: str,
+    max_new_tokens: int = 60,
+    beam_width: int = 8,
+    length_penalty: float = 0.7,
+    no_repeat_ngram_size: int = 3,
+    device: str = "cuda",
+    sft_mode: bool = False,
+    context_window: int = 2048,
+) -> str:
+    if sft_mode:
+        full_prompt = f"<|user|>\n{prompt}\n<|assistant|>\n"
+    else:
+        full_prompt = prompt
+    prompt_ids = tokenizer.encode(full_prompt, add_bos=True, add_eos=False)
+    prompt_len = len(prompt_ids)
+    stop_ids = build_stop_token_ids(tokenizer)
+    beams: List[Tuple[float, List[int]]] = [(0.0, list(prompt_ids))]
+    completed: List[Tuple[float, List[int]]] = []
+    for _step in range(max_new_tokens):
+        if not beams:
+            break
+        candidates: List[Tuple[float, List[int]]] = []
+        for beam_score, beam_ids in beams:
+            x = torch.tensor(
+                [beam_ids[-context_window:]], dtype=torch.long, device=device
+            )
+            with torch.no_grad():
+                logits, _, _, _ = model(x)
+            nl = logits[0, -1, :]
+            log_probs = F.log_softmax(nl, dim=-1)
+            gen_ids = beam_ids[prompt_len:]
+            if no_repeat_ngram_size > 1 and len(gen_ids) >= no_repeat_ngram_size - 1:
+                prefix = tuple(gen_ids[-(no_repeat_ngram_size - 1) :])
+                for i in range(len(gen_ids) - no_repeat_ngram_size + 1):
+                    if tuple(gen_ids[i : i + no_repeat_ngram_size - 1]) == prefix:
+                        log_probs[gen_ids[i + no_repeat_ngram_size - 1]] = float("-inf")
+            topk_lp, topk_ids = torch.topk(log_probs, beam_width)
+            for i in range(beam_width):
+                tid = topk_ids[i].item()
+                new_score = beam_score + topk_lp[i].item()
+                new_ids = beam_ids + [tid]
+                if tid in stop_ids:
+                    completed.append((new_score, new_ids))
+                else:
+                    candidates.append((new_score, new_ids))
+
+        def _norm_score(pair):
+            gen_len = max(1, len(pair[1]) - prompt_len)
+            return pair[0] / (gen_len**length_penalty)
+
+        candidates.sort(key=_norm_score, reverse=True)
+        beams = candidates[:beam_width]
+
+    pool = completed + beams
+    if not pool:
+        return ""
+
+    def _norm_score_final(pair):
+        gen_len = max(1, len(pair[1]) - prompt_len)
+        return pair[0] / (gen_len**length_penalty)
+
+    pool.sort(key=_norm_score_final, reverse=True)
+    best_ids = pool[0][1][prompt_len:]
+    text = tokenizer.decode(best_ids, skip_special=True)
+    nl_pos = text.find("\n")
+    if nl_pos > 5:
+        text = text[:nl_pos]
+    return text.strip()
+
+
+def generate(
+    model: TinyMemoryLM,
+    tokenizer: WordTokenizer,
+    prompt: str,
+    max_new_tokens: int = 256,
+    temperature: float = 0.8,
+    top_k: int = 40,
+    repetition_penalty: float = 1.0,
+    device: str = "cuda",
+    sft_mode: bool = False,
+    force_thought: bool = False,
+    stream: bool = True,
+    decode_mode: str = "legacy",
+    best_of: int = 3,
+    no_repeat_ngram_size: int = 3,
+    context_window: int = 2048,
+    beam_width: int = 8,
+    length_penalty: float = 0.7,
+) -> str:
+    if decode_mode == "beam":
+        text = generate_beam_search(
+            model=model,
+            tokenizer=tokenizer,
+            prompt=prompt,
+            max_new_tokens=max_new_tokens,
+            beam_width=beam_width,
+            length_penalty=length_penalty,
+            no_repeat_ngram_size=no_repeat_ngram_size,
+            device=device,
+            sft_mode=sft_mode,
+            context_window=context_window,
+        )
+        if stream:
+            print(text)
+        return text
+    if decode_mode == "legacy":
+        text, _, _, _ = generate_candidate(
+            model=model,
+            tokenizer=tokenizer,
+            prompt=prompt,
+            max_new_tokens=max_new_tokens,
+            temperature=temperature,
+            top_k=top_k,
+            repetition_penalty=repetition_penalty,
+            no_repeat_ngram_size=no_repeat_ngram_size,
+            device=device,
+            sft_mode=sft_mode,
+            force_thought=force_thought,
+            stream=stream,
+            context_window=context_window,
+        )
+        return text
+    candidates: List[Tuple[float, str, str, float]] = []
+    for _ in range(max(1, best_of)):
+        candidate_text, raw_text, avg_logprob, _ = generate_candidate(
+            model=model,
+            tokenizer=tokenizer,
+            prompt=prompt,
+            max_new_tokens=max_new_tokens,
+            temperature=temperature,
+            top_k=top_k,
+            repetition_penalty=repetition_penalty,
+            no_repeat_ngram_size=no_repeat_ngram_size,
+            device=device,
+            sft_mode=sft_mode,
+            force_thought=force_thought,
+            stream=False,
+            context_window=context_window,
+        )
+        score = score_candidate(prompt, raw_text, candidate_text, avg_logprob)
+        candidates.append((score, candidate_text, raw_text, avg_logprob))
+    best_score, best_text, _, _ = max(candidates, key=lambda item: item[0])
+    if stream:
+        print(best_text, end="", flush=True)
+        print()
+    return best_text
+
+
+# ---------------------------------------------------------------------------
+# Web server (from interactive.py)
+# ---------------------------------------------------------------------------
+
+ROOT = Path(__file__).resolve().parent
+if str(ROOT) not in sys.path:
+    sys.path.insert(0, str(ROOT))
+
+
+HF_ORG = "CompactAI"
+HF_API = "https://huggingface.co/api"
+CACHE_ROOT = Path.home() / ".cache" / "compactai_web"
+USER_AGENT = "Mozilla/5.0 CompactAI-Web"
+MODEL_CACHE: dict[tuple[str, str], dict[str, object]] = {}
+MODEL_CACHE_LOCK = threading.RLock()
+GENERATION_LOCK = threading.Lock()
+
+
+def request_json(url: str):
+    req = Request(url, headers={"User-Agent": USER_AGENT})
+    with urlopen(req, timeout=60) as response:
+        return json.loads(response.read().decode("utf-8"))
+
+
+def request_text(url: str) -> str:
+    req = Request(url, headers={"User-Agent": USER_AGENT})
+    with urlopen(req, timeout=60) as response:
+        return response.read().decode("utf-8", errors="replace")
+
+
+def download_file(url: str, destination: Path) -> None:
+    destination.parent.mkdir(parents=True, exist_ok=True)
+    temp_path = destination.with_suffix(destination.suffix + ".tmp")
+    req = Request(url, headers={"User-Agent": USER_AGENT})
+    with urlopen(req, timeout=120) as response, temp_path.open("wb") as handle:
+        shutil.copyfileobj(response, handle)
+    temp_path.replace(destination)
+
+
+def normalize_repo_id(raw_repo_id: str) -> str:
+    if not isinstance(raw_repo_id, str):
+        return ""
+    repo_id = raw_repo_id.strip()
+    if not repo_id:
+        return ""
+    try:
+        repo_id = unquote(repo_id)
+    except Exception:
+        pass
+    return (
+        repo_id.replace("https://huggingface.co/", "")
+        .replace("http://huggingface.co/", "")
+        .replace("api/models/", "")
+        .replace("models/", "")
+        .split("?", 1)[0]
+        .split("#", 1)[0]
+        .strip("/")
+    )
+
+
+def series_from_name(name: str) -> str | None:
+    lower = (name or "").lower()
+    if "haiku" in lower:
+        return "Haiku"
+    if "sonnet" in lower:
+        return "Sonnet"
+    if "opus" in lower:
+        return "Opus"
+    return None
+
+
+def encoded_repo_id(repo_id: str) -> str:
+    return "/".join(
+        quote(part, safe="") for part in normalize_repo_id(repo_id).split("/") if part
+    )
+
+
+def hf_file_url(repo_id: str, filename: str) -> str:
+    encoded_name = "/".join(
+        quote(part, safe="") for part in filename.split("/") if part
+    )
+    return (
+        f"https://huggingface.co/{encoded_repo_id(repo_id)}/resolve/main/{encoded_name}"
+    )
+
+
+def model_list() -> list[dict[str, object]]:
+    data = request_json(f"{HF_API}/models?author={quote(HF_ORG)}&full=true&limit=200")
+    models: list[dict[str, object]] = []
+    for item in data:
+        siblings = item.get("siblings") or []
+        filenames = [s.get("rfilename", "") for s in siblings if isinstance(s, dict)]
+        has_model = "model.pt" in filenames or "model/model.pt" in filenames
+        has_pretrain = "pretrain.pt" in filenames or "model/pretrain.pt" in filenames
+        has_tokenizer = (
+            "tokenizer.json" in filenames or "model/tokenizer.json" in filenames
+        )
+        if not has_model and not has_pretrain:
+            continue
+        name = (item.get("id") or "").split("/")[-1]
+        series = series_from_name(name)
+        if not series:
+            continue
+        models.append(
+            {
+                "id": item.get("id", ""),
+                "name": name,
+                "series": series,
+                "downloads": item.get("downloads", 0) or 0,
+                "likes": item.get("likes", 0) or 0,
+                "has_model": has_model,
+                "has_pretrain": has_pretrain,
+                "has_tokenizer": has_tokenizer,
+            }
+        )
+    return sorted(models, key=lambda entry: entry["downloads"], reverse=True)
+
+
+def model_details(repo_id: str) -> dict[str, object] | None:
+    normalized = normalize_repo_id(repo_id)
+    if not normalized:
+        return None
+    data = request_json(f"{HF_API}/models/{encoded_repo_id(normalized)}")
+    siblings = data.get("siblings") or []
+    files: dict[str, dict[str, float]] = {}
+    has_model = False
+    has_pretrain = False
+    for sibling in siblings:
+        if not isinstance(sibling, dict):
+            continue
+        filename = sibling.get("rfilename") or ""
+        if not filename:
+            continue
+        size_mb = round((sibling.get("size") or 0) / (1024 * 1024), 2)
+        files[filename] = {"size_mb": size_mb}
+        if filename.startswith("model/"):
+            files[filename.removeprefix("model/")] = {"size_mb": size_mb}
+        if filename in {"model.pt", "model/model.pt"}:
+            has_model = True
+        if filename in {"pretrain.pt", "model/pretrain.pt"}:
+            has_pretrain = True
+    readme_raw = ""
+    try:
+        readme_raw = request_text(
+            f"https://huggingface.co/{encoded_repo_id(normalized)}/raw/main/README.md"
+        )
+    except Exception:
+        readme_raw = ""
+    name = (data.get("id") or normalized).split("/")[-1]
+    return {
+        "id": normalized,
+        "name": name,
+        "series": series_from_name(name) or "Sonnet",
+        "downloads": data.get("downloads", 0) or 0,
+        "files": files,
+        "readme_raw": readme_raw,
+        "hf_model_id": normalized,
+        "has_model": has_model,
+        "has_pretrain": has_pretrain,
+    }
+
+
+def cache_dir(repo_id: str, model_type: str) -> Path:
+    return CACHE_ROOT / normalize_repo_id(repo_id).replace("/", "__") / model_type
+
+
+def artifact_candidates(model_type: str) -> list[str]:
+    return (
+        ["model/pretrain.pt", "pretrain.pt"]
+        if model_type == "pretrain"
+        else ["model/model.pt", "model.pt"]
+    )
+
+
+def ensure_artifact(repo_id: str, model_type: str, destination_name: str) -> Path:
+    normalized = normalize_repo_id(repo_id)
+    target = cache_dir(normalized, model_type) / destination_name
+    if target.exists():
+        return target
+    last_error: Exception | None = None
+    for candidate in (
+        artifact_candidates(model_type)
+        if destination_name.endswith(".pt")
+        else ["model/tokenizer.json", "tokenizer.json"]
+    ):
+        try:
+            download_file(hf_file_url(normalized, candidate), target)
+            return target
+        except Exception as exc:
+            last_error = exc
+    raise RuntimeError(
+        f"Unable to download {destination_name} for {normalized}: {last_error}"
+    )
+
+
+def series_config(series: str) -> dict[str, object]:
+    return MODEL_SERIES.get(series.lower(), MODEL_SERIES["sonnet"])
+
+
+def load_bundle(repo_id: str, model_type: str) -> dict[str, object]:
+    normalized = normalize_repo_id(repo_id)
+    details = model_details(normalized)
+    if not details:
+        raise RuntimeError("Model details are unavailable.")
+    series = str(details["series"])
+    key = (normalized, model_type)
+    with MODEL_CACHE_LOCK:
+        cached = MODEL_CACHE.get(key)
+        if cached:
+            return cached
+        bundle_dir = cache_dir(normalized, model_type)
+        bundle_dir.mkdir(parents=True, exist_ok=True)
+        model_path = bundle_dir / (
+            "pretrain.pt" if model_type == "pretrain" else "model.pt"
+        )
+        tokenizer_path = bundle_dir / "tokenizer.json"
+        if not model_path.exists():
+            ensure_artifact(normalized, model_type, model_path.name)
+        if not tokenizer_path.exists():
+            ensure_artifact(normalized, model_type, tokenizer_path.name)
+        tokenizer = WordTokenizer.load(tokenizer_path)
+        ckpt = torch.load(str(model_path), map_location="cpu", weights_only=False)
+        cfg = series_config(series)
+        vocab_size = int(ckpt.get("vocab_size", tokenizer.vocab_size))
+        state_dict = ckpt.get("model_state") or ckpt.get("state_dict") or ckpt
+        # Auto-detect new arch features from checkpoint weights
+        engram_dim = _detect_engram_dim(state_dict) or int(
+            cfg.get("engram_dim", model_config.engram_dim)
+        )
+        mhc_expansion = _detect_mhc_expansion(state_dict) or int(
+            cfg.get("mhc_expansion", model_config.mhc_expansion)
+        )
+        model = TinyMemoryLM(
+            vocab_size=vocab_size,
+            dim=int(cfg.get("dim", model_config.dim)),
+            n_unique_layers=int(
+                cfg.get("n_unique_layers", model_config.n_unique_layers)
+            ),
+            n_logical_layers=int(
+                cfg.get("n_logical_layers", model_config.n_logical_layers)
+            ),
+            n_heads=int(cfg.get("n_heads", model_config.n_heads)),
+            n_kv_heads=int(cfg.get("n_kv_heads", model_config.n_kv_heads)),
+            ffn_dim=int(cfg.get("ffn_dim", model_config.ffn_dim)),
+            dropout=float(cfg.get("dropout", model_config.dropout)),
+            mtp_horizons=tuple(
+                int(v) for v in cfg.get("mtp_horizons", model_config.mtp_horizons)
+            ),
+            grad_checkpoint=False,
+            sliding_window=int(
+                cfg.get("sliding_window_size", model_config.sliding_window_size)
+            ),
+            rope_fraction=float(
+                cfg.get("rope_fraction", model_config.rope_fraction)
+            ),
+            embed_scale=bool(
+                cfg.get("embed_scale", model_config.embed_scale)
+            ),
+            engram_dim=engram_dim,
+            engram_heads=int(cfg.get("engram_heads", model_config.engram_heads)),
+            engram_table_size=int(
+                cfg.get("engram_table_size", model_config.engram_table_size)
+            ),
+            engram_max_ngram=int(
+                cfg.get("engram_max_ngram", model_config.engram_max_ngram)
+            ),
+            mhc_expansion=mhc_expansion,
+        )
+        model.load_state_dict(state_dict, strict=False)
+        model.eval()
+        if tokenizer.vocab_size > vocab_size:
+            model.resize_token_embeddings(tokenizer.vocab_size)
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        model = model.to(device)
+        bundle = {
+            "repo_id": normalized,
+            "name": details["name"],
+            "series": series,
+            "type": model_type,
+            "model": model,
+            "tokenizer": tokenizer,
+            "device": device,
+            "model_path": str(model_path),
+            "tokenizer_path": str(tokenizer_path),
+            "downloads": details["downloads"],
+        }
+        MODEL_CACHE[key] = bundle
+        return bundle
+
+
+def ensure_port(start_port: int) -> int:
+    for port in range(start_port, start_port + 50):
+        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+            try:
+                sock.bind(("127.0.0.1", port))
+            except OSError:
+                continue
+            return port
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+        sock.bind(("127.0.0.1", 0))
+        return sock.getsockname()[1]
+
+
+def page_html() -> str:
+    return f"""<!doctype html>
+<html lang="en">
+<head>
+  <meta charset="utf-8">
+  <meta name="viewport" content="width=device-width, initial-scale=1">
+  <title>CompactAI Web</title>
+  <style>
+    :root {{
+      color-scheme: dark;
+      --bg: #050505;
+      --panel: #111111;
+      --panel-2: #161616;
+      --line: #262626;
+      --text: #f5f5f5;
+      --muted: #a3a3a3;
+      --accent: #d97706;
+      --accent-2: #b45309;
+      --soft: #1f1f1f;
+    }}
+    * {{ box-sizing: border-box; }}
+    body {{
+      margin: 0;
+      font-family: Geist, -apple-system, BlinkMacSystemFont, sans-serif;
+      background: var(--bg);
+      color: var(--text);
+      line-height: 1.5;
+    }}
+    a {{ color: inherit; }}
+    .wrap {{ max-width: 1120px; margin: 0 auto; padding: 28px 20px 40px; }}
+    .hero {{
+      display: flex;
+      justify-content: space-between;
+      align-items: end;
+      gap: 16px;
+      padding: 22px 0 28px;
+      border-bottom: 1px solid var(--line);
+      margin-bottom: 22px;
+    }}
+    h1 {{ margin: 0; font-size: clamp(2rem, 5vw, 3.5rem); letter-spacing: -0.04em; }}
+    .subtitle {{ margin: 10px 0 0; color: var(--muted); max-width: 58ch; }}
+    .grid {{
+      display: grid;
+      grid-template-columns: 1.1fr 1fr;
+      gap: 18px;
+    }}
+    .panel {{
+      background: var(--panel);
+      border: 1px solid var(--line);
+      border-radius: 18px;
+      padding: 18px;
+    }}
+    .panel h2 {{ margin: 0 0 12px; font-size: 15px; letter-spacing: 0.02em; text-transform: uppercase; color: var(--muted); }}
+    .row {{ display: flex; gap: 10px; flex-wrap: wrap; }}
+    select, textarea, input {{
+      width: 100%;
+      background: var(--panel-2);
+      color: var(--text);
+      border: 1px solid var(--line);
+      border-radius: 12px;
+      padding: 12px 14px;
+      font: inherit;
+      outline: none;
+    }}
+    textarea {{ min-height: 170px; resize: vertical; }}
+    select {{ appearance: none; }}
+    .choice {{
+      flex: 1 1 150px;
+      display: flex;
+      align-items: center;
+      gap: 10px;
+      padding: 10px 12px;
+      border: 1px solid var(--line);
+      border-radius: 12px;
+      background: var(--panel-2);
+      cursor: pointer;
+    }}
+    .choice input {{ width: auto; }}
+    .btns {{ display: flex; flex-wrap: wrap; gap: 10px; }}
+    button {{
+      border: 1px solid var(--line);
+      border-radius: 12px;
+      padding: 11px 14px;
+      background: var(--soft);
+      color: var(--text);
+      font: inherit;
+      cursor: pointer;
+      transition: transform 0.15s ease, border-color 0.15s ease, background 0.15s ease;
+    }}
+    button:hover {{ transform: translateY(-1px); border-color: #3a3a3a; }}
+    .primary {{ background: var(--accent); border-color: var(--accent); color: #fff; }}
+    .primary:hover {{ background: var(--accent-2); border-color: var(--accent-2); }}
+    .status {{
+      margin-top: 12px;
+      color: var(--muted);
+      font-size: 13px;
+      min-height: 1.4em;
+    }}
+    .output {{
+      white-space: pre-wrap;
+      background: #0b0b0b;
+      border: 1px solid var(--line);
+      border-radius: 16px;
+      min-height: 280px;
+      padding: 16px;
+      color: #e7e5e4;
+      overflow: auto;
+    }}
+    .meta {{
+      display: flex;
+      flex-wrap: wrap;
+      gap: 8px;
+      margin-top: 8px;
+    }}
+    .chip {{
+      display: inline-flex;
+      align-items: center;
+      gap: 6px;
+      padding: 6px 10px;
+      border-radius: 999px;
+      border: 1px solid var(--line);
+      background: var(--panel-2);
+      font-size: 12px;
+      color: var(--muted);
+    }}
+    .code {{
+      margin-top: 14px;
+      padding: 12px 14px;
+      border-radius: 12px;
+      border: 1px solid var(--line);
+      background: #0b0b0b;
+      font-family: ui-monospace, SFMono-Regular, Menlo, Monaco, Consolas, monospace;
+      font-size: 13px;
+      overflow-x: auto;
+    }}
+    @media (max-width: 900px) {{
+      .grid {{ grid-template-columns: 1fr; }}
+      .hero {{ align-items: start; flex-direction: column; }}
+    }}
+  </style>
+</head>
+<body>
+  <div class="wrap">
+    <div class="hero">
+      <div>
+        <h1>CompactAI Web</h1>
+        <p class="subtitle">Pull a model from Hugging Face, keep it cached locally, and chat in the browser.</p>
+      </div>
+      <div class="meta">
+        <span class="chip">Hugging Face: CompactAI</span>
+        <span class="chip">pip install -r requirements.txt</span>
+        <span class="chip">Local inference</span>
+      </div>
+    </div>
+
+    <div class="grid">
+      <section class="panel">
+        <h2>Model</h2>
+        <select id="modelSelect"></select>
+        <div class="row" style="margin-top: 10px;">
+          <label class="choice"><input type="radio" name="type" value="model" checked> Instruct / final</label>
+          <label class="choice"><input type="radio" name="type" value="pretrain"> Pretrain</label>
+        </div>
+        <div class="btns" style="margin-top: 12px;">
+          <button id="downloadBtn">Download</button>
+          <button id="refreshBtn">Refresh models</button>
+        </div>
+        <div class="status" id="modelStatus">Loading model list…</div>
+        <div class="code">python3 interactive_web.py</div>
+      </section>
+
+      <section class="panel">
+        <h2>Prompt</h2>
+        <textarea id="prompt" placeholder="Ask something…"></textarea>
+        <div class="row" style="margin-top: 10px;">
+          <input id="temperature" type="number" min="0.1" max="2" step="0.05" value="0.8" style="flex: 1 1 120px;">
+          <input id="topK" type="number" min="1" max="100" step="1" value="40" style="flex: 1 1 120px;">
+          <input id="maxTokens" type="number" min="16" max="2048" step="16" value="256" style="flex: 1 1 120px;">
+        </div>
+        <div class="btns" style="margin-top: 12px;">
+          <button id="generateBtn" class="primary">Generate</button>
+        </div>
+        <div class="status" id="genStatus"></div>
+      </section>
+    </div>
+
+    <section class="panel" style="margin-top: 18px;">
+      <h2>Response</h2>
+      <div id="output" class="output"></div>
+    </section>
+  </div>
+
+  <script>
+    const modelSelect = document.getElementById('modelSelect');
+    const modelStatus = document.getElementById('modelStatus');
+    const genStatus = document.getElementById('genStatus');
+    const output = document.getElementById('output');
+    const promptBox = document.getElementById('prompt');
+
+    async function api(path, body) {{
+      const response = await fetch(path, {{
+        method: body ? 'POST' : 'GET',
+        headers: body ? {{ 'Content-Type': 'application/json' }} : undefined,
+        body: body ? JSON.stringify(body) : undefined,
+      }});
+      return response.json();
+    }}
+
+    function currentType() {{
+      return document.querySelector('input[name="type"]:checked').value;
+    }}
+
+    function currentModelId() {{
+      return modelSelect.value;
+    }}
+
+    function setModels(models) {{
+      modelSelect.innerHTML = '';
+      for (const model of models) {{
+        const option = document.createElement('option');
+        option.value = model.id;
+        option.textContent = `${{model.name}} • ${{model.series}}`;
+        modelSelect.appendChild(option);
+      }}
+      if (models.length === 0) {{
+        const option = document.createElement('option');
+        option.value = '';
+        option.textContent = 'No CompactAI models found';
+        modelSelect.appendChild(option);
+      }}
+    }}
+
+    async function refreshModels() {{
+      modelStatus.textContent = 'Loading model list…';
+      try {{
+        const models = await api('/api/models');
+        setModels(models);
+        modelStatus.textContent = models.length ? `${{models.length}} models available from CompactAI` : 'No compatible models found.';
+      }} catch (error) {{
+        modelStatus.textContent = 'Failed to load model list.';
+      }}
+    }}
+
+    async function ensureModel() {{
+      const modelId = currentModelId();
+      if (!modelId) {{
+        modelStatus.textContent = 'Pick a model first.';
+        return null;
+      }}
+      modelStatus.textContent = 'Downloading model files…';
+      const result = await api('/api/ensure', {{ modelId, type: currentType() }});
+      if (!result.success) {{
+        modelStatus.textContent = result.error || 'Download failed.';
+        return null;
+      }}
+      modelStatus.textContent = `${{result.name}} ready on ${{result.series}}`;
+      return result;
+    }}
+
+    async function generate() {{
+      output.textContent = '';
+      genStatus.textContent = '';
+      const modelId = currentModelId();
+      const prompt = promptBox.value.trim();
+      if (!modelId) {{
+        genStatus.textContent = 'Pick a model first.';
+        return;
+      }}
+      if (!prompt) {{
+        genStatus.textContent = 'Enter a prompt first.';
+        return;
+      }}
+      genStatus.textContent = 'Preparing model…';
+      const result = await api('/api/generate', {{
+        modelId,
+        type: currentType(),
+        prompt,
+        temperature: Number(document.getElementById('temperature').value || 0.8),
+        top_k: Number(document.getElementById('topK').value || 40),
+        max_new_tokens: Number(document.getElementById('maxTokens').value || 256),
+      }});
+      if (!result.success) {{
+        genStatus.textContent = result.error || 'Generation failed.';
+        return;
+      }}
+      output.textContent = result.text || '';
+      genStatus.textContent = 'Done.';
+    }}
+
+    document.getElementById('refreshBtn').addEventListener('click', refreshModels);
+    document.getElementById('downloadBtn').addEventListener('click', ensureModel);
+    document.getElementById('generateBtn').addEventListener('click', generate);
+    promptBox.addEventListener('keydown', (event) => {{
+      if (event.key === 'Enter' && (event.ctrlKey || event.metaKey)) {{
+        event.preventDefault();
+        generate();
+      }}
+    }});
+
+    refreshModels();
+  </script>
+</body>
+</html>"""
+
+
+class Handler(BaseHTTPRequestHandler):
+    def _send_json(self, payload, status=200):
+        body = json.dumps(payload).encode("utf-8")
+        self.send_response(status)
+        self.send_header("Content-Type", "application/json; charset=utf-8")
+        self.send_header("Content-Length", str(len(body)))
+        self.send_header("Cache-Control", "no-store")
+        self.end_headers()
+        self.wfile.write(body)
+
+    def _send_html(self, payload: str, status=200):
+        body = payload.encode("utf-8")
+        self.send_response(status)
+        self.send_header("Content-Type", "text/html; charset=utf-8")
+        self.send_header("Content-Length", str(len(body)))
+        self.send_header("Cache-Control", "no-store")
+        self.end_headers()
+        self.wfile.write(body)
+
+    def do_GET(self):
+        parsed = urlparse(self.path)
+        if parsed.path in {"/", "/index.html"}:
+            self._send_html(page_html())
+            return
+        if parsed.path == "/api/models":
+            try:
+                self._send_json(model_list())
+            except Exception as exc:
+                self._send_json({"success": False, "error": str(exc)}, 500)
+            return
+        if parsed.path.startswith("/api/models/"):
+            repo_id = normalize_repo_id(parsed.path.removeprefix("/api/models/"))
+            try:
+                details = model_details(repo_id)
+                if not details:
+                    self._send_json(
+                        {"success": False, "error": "Model not found."}, 404
+                    )
+                else:
+                    self._send_json(details)
+            except Exception as exc:
+                self._send_json({"success": False, "error": str(exc)}, 500)
+            return
+        self._send_json({"success": False, "error": "Not found."}, 404)
+
+    def do_POST(self):
+        parsed = urlparse(self.path)
+        length = int(self.headers.get("Content-Length", "0") or "0")
+        raw = self.rfile.read(length).decode("utf-8") if length else "{}"
+        try:
+            payload = json.loads(raw or "{}")
+        except Exception:
+            payload = {}
+        if parsed.path == "/api/ensure":
+            try:
+                repo_id = normalize_repo_id(payload.get("modelId", ""))
+                model_type = payload.get("type", "model")
+                if not repo_id:
+                    self._send_json(
+                        {"success": False, "error": "Missing model ID."}, 400
+                    )
+                    return
+                details = model_details(repo_id)
+                if not details:
+                    self._send_json(
+                        {"success": False, "error": "Model not found."}, 404
+                    )
+                    return
+                bundle = load_bundle(repo_id, model_type)
+                self._send_json(
+                    {
+                        "success": True,
+                        "id": bundle["repo_id"],
+                        "name": bundle["name"],
+                        "series": bundle["series"],
+                        "type": bundle["type"],
+                    }
+                )
+            except Exception as exc:
+                self._send_json({"success": False, "error": str(exc)}, 500)
+            return
+        if parsed.path == "/api/generate":
+            try:
+                repo_id = normalize_repo_id(payload.get("modelId", ""))
+                model_type = payload.get("type", "model")
+                prompt = str(payload.get("prompt", ""))
+                if not repo_id:
+                    self._send_json(
+                        {"success": False, "error": "Missing model ID."}, 400
+                    )
+                    return
+                bundle = load_bundle(repo_id, model_type)
+                with GENERATION_LOCK:
+                    text = generate(
+                        model=bundle["model"],
+                        tokenizer=bundle["tokenizer"],
+                        prompt=prompt,
+                        max_new_tokens=int(payload.get("max_new_tokens", 256)),
+                        temperature=float(payload.get("temperature", 0.8)),
+                        top_k=int(payload.get("top_k", 40)),
+                        repetition_penalty=float(
+                            payload.get("repetition_penalty", 1.0)
+                        ),
+                        device=str(bundle["device"]),
+                        sft_mode=model_type != "pretrain",
+                        force_thought=bool(payload.get("force_thought", False)),
+                        stream=False,
+                        decode_mode=str(payload.get("decode_mode", "legacy")),
+                        best_of=int(payload.get("best_of", 3)),
+                        no_repeat_ngram_size=int(
+                            payload.get("no_repeat_ngram_size", 3)
+                        ),
+                        context_window=int(payload.get("context_window", 2048)),
+                        beam_width=int(payload.get("beam_width", 8)),
+                        length_penalty=float(payload.get("length_penalty", 0.7)),
+                    )
+                self._send_json(
+                    {
+                        "success": True,
+                        "text": text,
+                        "name": bundle["name"],
+                        "series": bundle["series"],
+                    }
+                )
+            except Exception as exc:
+                self._send_json({"success": False, "error": str(exc)}, 500)
+            return
+        self._send_json({"success": False, "error": "Not found."}, 404)
+
+    def log_message(self, format, *args):
+        return
+
+
+def main():
+    CACHE_ROOT.mkdir(parents=True, exist_ok=True)
+    port = ensure_port(int(os.environ.get("PORT", "7860")))
+    server = ThreadingHTTPServer(("127.0.0.1", port), Handler)
+    url = f"http://127.0.0.1:{port}"
+    print(url, flush=True)
+    try:
+        webbrowser.open(url)
+    except Exception:
+        pass
+    try:
+        server.serve_forever()
+    except KeyboardInterrupt:
+        pass
+    finally:
+        server.server_close()
+
+
+if __name__ == "__main__":
+    main()

requirements.txt

@@ -0,0 +1 @@
+torch>=2.0.0