Update downloads/interactive.py

downloads/interactive.py

@@ -25,10 +25,12 @@ import torch.nn as nn
 import torch.nn.functional as F
 from torch.utils.checkpoint import checkpoint
 
+
 HUGGINGFACE_MODELS = {
     "TMLM-Haiku-1": "CompactAI-O/TMLM-Haiku-1",
     "TMLM-Haiku-1.3": "CompactAI-O/TMLM-Haiku-1.3",
     "TMLM-Haiku-2": "CompactAI-O/TMLM-Haiku-2",
+    "Glint-1": "CompactAI-O/Glint-1",
 }
 
 
@@ -79,6 +81,15 @@ MODEL_SERIES = {
         "engram_table_size": 64,
         "engram_max_ngram": 2,
         "mhc_expansion": 2,
+        "sleep_gate_cap": 0,
+        "sleep_gate_heads": 4,
+        "latent_think_layers": 0,
+        "prelude_layers": 0,
+        "coda_layers": 0,
+        "recurrent_loops": 0,
+        "recurrent_act_threshold": 0.9,
+        "recurrent_lora_rank": 0,
+        "recurrent_loop_embed_dim": 0,
     },
     "sonnet": {
         "dim": 1024,
@@ -95,6 +106,15 @@ MODEL_SERIES = {
         "engram_table_size": 4096,
         "engram_max_ngram": 2,
         "mhc_expansion": 2,
+        "sleep_gate_cap": 0,
+        "sleep_gate_heads": 8,
+        "latent_think_layers": 0,
+        "prelude_layers": 0,
+        "coda_layers": 0,
+        "recurrent_loops": 0,
+        "recurrent_act_threshold": 0.99,
+        "recurrent_lora_rank": 0,
+        "recurrent_loop_embed_dim": 0,
     },
     "opus": {
         "dim": 1536,
@@ -111,6 +131,15 @@ MODEL_SERIES = {
         "engram_table_size": 8192,
         "engram_max_ngram": 2,
         "mhc_expansion": 4,
+        "sleep_gate_cap": 0,
+        "sleep_gate_heads": 8,
+        "latent_think_layers": 0,
+        "prelude_layers": 0,
+        "coda_layers": 0,
+        "recurrent_loops": 0,
+        "recurrent_act_threshold": 0.99,
+        "recurrent_lora_rank": 0,
+        "recurrent_loop_embed_dim": 0,
     },
 }
 
@@ -422,6 +451,68 @@ class SwiGLU(nn.Module):
         return out
 
 
+def loop_index_embedding(h: torch.Tensor, loop_t: int, loop_dim: int, theta: float = 10000.0) -> torch.Tensor:
+    if loop_dim <= 0:
+        return h
+    loop_dim = min(loop_dim, h.shape[-1])
+    if loop_dim % 2 == 1:
+        loop_dim -= 1
+    if loop_dim <= 0:
+        return h
+    inv_freq = 1.0 / (theta ** (torch.arange(0, loop_dim, 2, device=h.device, dtype=h.dtype) / loop_dim))
+    phase = torch.tensor(float(loop_t), device=h.device, dtype=h.dtype) * inv_freq
+    loop_embed = torch.cat([phase.sin(), phase.cos()], dim=0).view(1, 1, loop_dim)
+    out = h.clone()
+    out[..., :loop_dim] = out[..., :loop_dim] + loop_embed
+    return out
+
+
+class DepthLoRAAdapter(nn.Module):
+    def __init__(self, dim: int, rank: int, max_loops: int) -> None:
+        super().__init__()
+        self.rank = max(0, rank)
+        if self.rank <= 0:
+            self.down = None
+            self.B = None
+            self.scale = None
+            return
+        self.down = nn.Linear(dim, self.rank, bias=False)
+        self.B = nn.Parameter(torch.randn(self.rank, dim) * 0.02)
+        self.scale = nn.Embedding(max(1, max_loops), self.rank)
+        nn.init.zeros_(self.scale.weight)
+
+    def forward(self, x: torch.Tensor, loop_t: int) -> torch.Tensor:
+        if self.rank <= 0 or self.down is None or self.B is None or self.scale is None:
+            return torch.zeros_like(x)
+        t_idx = min(loop_t, self.scale.num_embeddings - 1)
+        scale = self.scale(torch.tensor(t_idx, device=x.device))
+        return (self.down(x) * scale) @ self.B
+
+
+class StableRecurrentInjection(nn.Module):
+    def __init__(self, dim: int) -> None:
+        super().__init__()
+        self.log_A = nn.Parameter(torch.full((dim,), -2.0))
+        self.log_dt = nn.Parameter(torch.full((dim,), -2.0))
+        self.input_gate = nn.Parameter(torch.zeros(dim))
+
+    def forward(self, h: torch.Tensor, e: torch.Tensor, transformer_out: torch.Tensor) -> torch.Tensor:
+        A = torch.exp(-torch.exp((self.log_dt + self.log_A).clamp(-20, 20))).view(1, 1, -1)
+        B = torch.sigmoid(self.input_gate).view(1, 1, -1)
+        return A * h + B * e + transformer_out
+
+
+class AdaptiveHalting(nn.Module):
+    def __init__(self, dim: int) -> None:
+        super().__init__()
+        self.halt = nn.Linear(dim, 1, bias=True)
+        nn.init.zeros_(self.halt.weight)
+        nn.init.constant_(self.halt.bias, -2.0)
+
+    def forward(self, h: torch.Tensor) -> torch.Tensor:
+        return torch.sigmoid(self.halt(h)).squeeze(-1)
+
+
 class EngramBlock(nn.Module):
     """DeepSeek Engram: conditional memory via O(1) hashed N-gram lookup.
 
@@ -566,6 +657,115 @@ class EngramBlock(nn.Module):
         return gate * value
 
 
+class SleepGate(nn.Module):
+    """Persistent memory + periodic consolidation gate."""
+
+    def __init__(
+        self,
+        dim: int,
+        cap: int = 128,
+        n_heads: int = 4,
+        retention_enabled: bool = True,
+        retention_hidden: int = 0,
+    ) -> None:
+        super().__init__()
+        self.dim = dim
+        self.cap = cap
+        self.n_heads = n_heads
+        self.head_dim = dim // n_heads
+        self.scale = self.head_dim ** -0.5
+        self.retention_enabled = retention_enabled
+
+        self.register_buffer("mem_emb", torch.zeros(cap, dim, dtype=torch.bfloat16))
+        self.register_buffer("mem_age", torch.zeros(cap, dtype=torch.long))
+        self.register_buffer("mem_beta", torch.ones(cap, dtype=torch.float32))
+        self.register_buffer("mem_count", torch.zeros((), dtype=torch.long))
+        self.register_buffer("mem_head", torch.zeros((), dtype=torch.long))
+        self.register_buffer("global_step", torch.zeros((), dtype=torch.long))
+
+        self.q_proj = nn.Linear(dim, dim, bias=False)
+        self.k_proj = nn.Linear(dim, dim, bias=False)
+        self.v_proj = nn.Linear(dim, dim, bias=False)
+        self.o_proj = nn.Linear(dim, dim, bias=False)
+        nn.init.zeros_(self.o_proj.weight)
+        self.gate_scale = nn.Parameter(torch.zeros(()))
+
+        if retention_enabled:
+            if retention_hidden > 0:
+                self.retention_gate: Optional[nn.Module] = nn.Sequential(
+                    nn.Linear(dim, retention_hidden, bias=False),
+                    nn.GELU(),
+                    nn.Linear(retention_hidden, 1, bias=True),
+                )
+                nn.init.constant_(self.retention_gate[-1].bias, 2.2)
+            else:
+                self.retention_gate = nn.Linear(dim, 1, bias=True)
+                nn.init.constant_(self.retention_gate.bias, 2.2)
+        else:
+            self.retention_gate = None
+
+        self._last_beta: Optional[torch.Tensor] = None
+
+    def write(self, hidden: torch.Tensor) -> None:
+        B, T, _ = hidden.shape
+        tail_full = hidden[:, max(0, T - 16):, :].float().mean(dim=1)
+        if self.retention_gate is not None:
+            beta_live = torch.sigmoid(self.retention_gate(tail_full).squeeze(-1))
+            self._last_beta = beta_live if self.training else None
+            beta_store = beta_live.detach().float()
+        else:
+            self._last_beta = None
+            beta_store = torch.ones(B, device=hidden.device, dtype=torch.float32)
+        tail = tail_full.to(self.mem_emb.dtype).detach()
+        with torch.no_grad():
+            head = int(self.mem_head.item())
+            count = int(self.mem_count.item())
+            step = int(self.global_step.item())
+            for b in range(B):
+                self.mem_emb[head] = tail[b]
+                self.mem_age[head] = step
+                self.mem_beta[head] = beta_store[b]
+                head = (head + 1) % self.cap
+                if count < self.cap:
+                    count += 1
+            self.mem_head.fill_(head)
+            self.mem_count.fill_(count)
+
+    def read(self, x: torch.Tensor) -> torch.Tensor:
+        count = int(self.mem_count.item())
+        if count == 0:
+            return torch.zeros_like(x)
+        B, T, D = x.shape
+        mem = self.mem_emb[:count].clone().to(x.dtype)
+        q = self.q_proj(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
+        k = self.k_proj(mem).view(count, self.n_heads, self.head_dim).transpose(0, 1)
+        v = self.v_proj(mem).view(count, self.n_heads, self.head_dim).transpose(0, 1)
+        attn = torch.einsum("bhtd,hmd->bhtm", q, k) * self.scale
+        attn = F.softmax(attn, dim=-1)
+        if self.retention_enabled:
+            step = int(self.global_step.item())
+            ages = self.mem_age[:count].to(x.device)
+            delta = (step - ages).clamp(min=0).to(x.dtype)
+            betas = self.mem_beta[:count].to(x.dtype).clamp(min=1e-6, max=1.0)
+            weights = betas.pow(delta)
+            attn = attn * weights.view(1, 1, 1, count)
+            attn = attn / attn.sum(dim=-1, keepdim=True).clamp_min(1e-9)
+        out = torch.einsum("bhtm,hmd->bhtd", attn, v)
+        out = out.transpose(1, 2).contiguous().view(B, T, D)
+        out = self.o_proj(out)
+        return torch.sigmoid(self.gate_scale) * out
+
+    @torch.no_grad()
+    def reset(self) -> None:
+        self.mem_emb.zero_()
+        self.mem_age.zero_()
+        self.mem_beta.fill_(1.0)
+        self.mem_count.zero_()
+        self.mem_head.zero_()
+        self.global_step.zero_()
+        self._last_beta = None
+
+
 def _sinkhorn_knopp(logits: torch.Tensor, n_iters: int = 7) -> torch.Tensor:
     M = torch.exp(logits.clamp(-10, 10))
     for _ in range(n_iters):
@@ -733,6 +933,85 @@ class TransformerBlock(nn.Module):
         return x, new_kv
 
 
+class RecurrentDepthBlock(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,
+        n_loops: int,
+        act_threshold: float,
+        lora_rank: int,
+        loop_embed_dim: int,
+    ) -> None:
+        super().__init__()
+        self.n_loops = max(1, n_loops)
+        self.act_threshold = act_threshold
+        self.loop_embed_dim = max(0, loop_embed_dim)
+        self.norm = RMSNorm(dim)
+        self.block = 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=0, mhc_expansion=1,
+        )
+        self.injection = StableRecurrentInjection(dim)
+        self.act = AdaptiveHalting(dim)
+        self.lora = DepthLoRAAdapter(dim, lora_rank, self.n_loops)
+
+    def forward(
+        self,
+        h: torch.Tensor,
+        e: torch.Tensor,
+        token_ids: Optional[torch.Tensor] = None,
+        past_key_values: Optional[List[Optional[Tuple[torch.Tensor, torch.Tensor]]]] = None,
+        use_cache: bool = False,
+        n_loops: Optional[int] = None,
+    ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor], Optional[List[Tuple[torch.Tensor, torch.Tensor]]]]:
+        loops = max(1, n_loops or self.n_loops)
+        B, T, _ = h.shape
+        halted = torch.zeros(B, T, device=h.device, dtype=torch.bool)
+        cumulative_p = torch.zeros(B, T, device=h.device, dtype=h.dtype)
+        output = torch.zeros_like(h)
+        new_past: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = [] if use_cache else None
+        current = h
+        final_halt = None
+
+        for t in range(loops):
+            h_loop = loop_index_embedding(current, t, self.loop_embed_dim)
+            combined = self.norm(h_loop + e)
+            past_kv = None
+            if past_key_values is not None and t < len(past_key_values):
+                past_kv = past_key_values[t]
+            trans_out, layer_kv = self.block(combined, is_global=True, past_kv=past_kv, use_cache=use_cache, token_ids=token_ids)
+            trans_out = trans_out + self.lora(trans_out, t)
+            next_h = self.injection(current, e, trans_out)
+            p = self.act(next_h)
+            p = p * (~halted).to(p.dtype)
+            final_halt = p
+            should_halt = (~halted) & ((cumulative_p + p) >= self.act_threshold)
+            update_weight = torch.where(should_halt, (1.0 - cumulative_p).clamp(min=0.0), p)
+            output = output + next_h * update_weight.unsqueeze(-1)
+            cumulative_p = cumulative_p + update_weight
+            current = torch.where(halted.unsqueeze(-1), current, next_h)
+            halted = halted | should_halt
+            if new_past is not None:
+                new_past.append(layer_kv)
+            if not use_cache and bool(halted.all()):
+                break
+
+        remainder = (1.0 - cumulative_p).clamp(min=0.0)
+        output = output + current * remainder.unsqueeze(-1)
+        aux: Dict[str, torch.Tensor] = {}
+        if final_halt is not None:
+            aux["recurrent_halt_mean"] = final_halt.mean()
+        return output, aux, new_past
+
+
 class TinyMemoryLM(nn.Module):
     def __init__(
         self,
@@ -754,6 +1033,17 @@ class TinyMemoryLM(nn.Module):
         engram_table_size: int = 8192,
         engram_max_ngram: int = 3,
         mhc_expansion: int = 1,
+        sleep_gate_cap: int = 0,
+        sleep_gate_heads: int = 4,
+        sleep_retention_enabled: bool = True,
+        sleep_retention_hidden: int = 0,
+        latent_think_layers: int = 0,
+        prelude_layers: int = 0,
+        coda_layers: int = 0,
+        recurrent_loops: int = 0,
+        recurrent_act_threshold: float = 0.99,
+        recurrent_lora_rank: int = 0,
+        recurrent_loop_embed_dim: int = 0,
     ) -> None:
         super().__init__()
         self.dim = dim
@@ -766,29 +1056,45 @@ class TinyMemoryLM(nn.Module):
         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(
-            [
+        self.use_recurrent_depth = recurrent_loops > 0
+        self.prelude_layers = max(0, prelude_layers)
+        self.coda_layers = max(0, coda_layers)
+        self.recurrent_loops = max(0, recurrent_loops)
+
+        self.blocks: Optional[nn.ModuleList] = None
+        self.prelude: Optional[nn.ModuleList] = None
+        self.recurrent: Optional[RecurrentDepthBlock] = None
+        self.coda: Optional[nn.ModuleList] = None
+
+        def _make_blocks(n: int) -> nn.ModuleList:
+            return 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,
+                    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)
-            ]
-        )
+                for _ in range(n)
+            ])
+
+        if self.use_recurrent_depth:
+            if self.prelude_layers > 0:
+                self.prelude = _make_blocks(self.prelude_layers)
+            self.recurrent = RecurrentDepthBlock(
+                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, n_loops=self.recurrent_loops,
+                act_threshold=recurrent_act_threshold, lora_rank=recurrent_lora_rank,
+                loop_embed_dim=recurrent_loop_embed_dim or max(2, dim // 8),
+            )
+            if self.coda_layers > 0:
+                self.coda = _make_blocks(self.coda_layers)
+        else:
+            self.blocks = _make_blocks(max(1, n_unique_layers))
+
         self.norm = RMSNorm(dim)
 
         self.mtp_horizons = sorted({int(h) for h in mtp_horizons if int(h) > 1})
@@ -799,10 +1105,37 @@ class TinyMemoryLM(nn.Module):
             {str(h): RMSNorm(dim) for h in self.mtp_horizons}
         )
 
-        res_scale = (2 * n_unique_layers) ** -0.5
-        for block in self.blocks:
-            block.attn.wo.weight.data.mul_(res_scale)
-            block.ffn.down.weight.data.mul_(res_scale)
+        res_scale = (2 * max(1, n_logical_layers)) ** -0.5
+        for group in (self.blocks, self.prelude, self.coda):
+            if group is None:
+                continue
+            for block in group:
+                block.attn.wo.weight.data.mul_(res_scale)
+                block.ffn.down.weight.data.mul_(res_scale)
+        if self.recurrent is not None:
+            self.recurrent.block.attn.wo.weight.data.mul_(res_scale)
+            self.recurrent.block.ffn.down.weight.data.mul_(res_scale)
+
+        self.sleep_gate: Optional[SleepGate] = None
+        if sleep_gate_cap > 0:
+            self.sleep_gate = SleepGate(
+                dim=dim, cap=sleep_gate_cap, n_heads=sleep_gate_heads,
+                retention_enabled=sleep_retention_enabled,
+                retention_hidden=sleep_retention_hidden,
+            )
+
+        self.think_blocks: Optional[nn.ModuleList] = None
+        self.think_norm: Optional[RMSNorm] = None
+        if latent_think_layers > 0:
+            self.think_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=0.0, sliding_window=2048,
+                    rope_fraction=rope_fraction, engram_dim=0, mhc_expansion=1,
+                )
+                for _ in range(latent_think_layers)
+            ])
+            self.think_norm = RMSNorm(dim)
 
     def resize_token_embeddings(self, new_vocab_size: int) -> None:
         old_vocab_size = self.embed_tokens.num_embeddings
@@ -810,12 +1143,8 @@ class TinyMemoryLM(nn.Module):
             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.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))
@@ -824,62 +1153,74 @@ class TinyMemoryLM(nn.Module):
         self.embed_tokens.weight.data[:copy_size] = old_embed_weight[:copy_size]
 
     def _build_logical_layers(self) -> List[Tuple[nn.Module, int]]:
-        logical = []
+        if self.blocks is None:
+            return []
         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
+        return [(block, i) for i, block in enumerate(full_sequence[: self.n_logical_layers])]
 
     def forward(
         self,
         ids: torch.Tensor,
         use_cache: bool = False,
-        past_key_values: Optional[
-            List[Optional[Tuple[torch.Tensor, torch.Tensor]]]
-        ] = None,
+        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]]],
-    ]:
+    ) -> Tuple[torch.Tensor, Dict[int, torch.Tensor], Dict[str, 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
-        token_ids = ids
-
-        logical_layers = self._build_logical_layers()
-        new_past_key_values: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = (
-            [] if use_cache else None
-        )
-
-        last_logical_idx = len(logical_layers) - 1
-        for layer_idx, (block, logical_idx) in enumerate(logical_layers):
-            is_global = logical_idx % 2 == 0 or layer_idx == last_logical_idx
-            past_kv = (
-                past_key_values[layer_idx]
-                if past_key_values is not None and layer_idx < len(past_key_values)
-                else None
+        new_past_key_values: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = [] if use_cache else None
+        aux: Dict[str, torch.Tensor] = {}
+
+        if self.use_recurrent_depth:
+            offset = 0
+            if self.prelude is not None:
+                for block in self.prelude:
+                    past_kv = past_key_values[offset] if past_key_values is not None and offset < len(past_key_values) else None
+                    x, layer_kv = block(x, is_global=True, past_kv=past_kv, use_cache=use_cache, token_ids=ids)
+                    if new_past_key_values is not None:
+                        new_past_key_values.append(layer_kv)
+                    offset += 1
+            encoded = x
+            recurrent_past = past_key_values[offset: offset + self.recurrent_loops] if past_key_values is not None else None
+            x, recurrent_aux, recurrent_kv = self.recurrent(
+                x, encoded, token_ids=ids, past_key_values=recurrent_past, use_cache=use_cache,
             )
+            aux.update(recurrent_aux)
+            if new_past_key_values is not None and recurrent_kv is not None:
+                new_past_key_values.extend(recurrent_kv)
+            offset += self.recurrent_loops
+            if self.coda is not None:
+                for block in self.coda:
+                    past_kv = past_key_values[offset] if past_key_values is not None and offset < len(past_key_values) else None
+                    x, layer_kv = block(x, is_global=True, past_kv=past_kv, use_cache=use_cache, token_ids=ids)
+                    if new_past_key_values is not None:
+                        new_past_key_values.append(layer_kv)
+                    offset += 1
+        else:
+            logical_layers = self._build_logical_layers()
+            last_logical_idx = len(logical_layers) - 1
+            for layer_idx, (block, logical_idx) in enumerate(logical_layers):
+                is_global = logical_idx % 2 == 0 or layer_idx == last_logical_idx
+                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)
 
-            if self.grad_checkpoint and self.training and not use_cache:
-                x, layer_kv = checkpoint(
-                    block,
-                    x,
-                    is_global,
-                    past_kv,
-                    use_cache,
-                    token_ids,
-                    use_reentrant=True,
-                )
-            else:
-                x, layer_kv = block(x, is_global, past_kv, use_cache, token_ids)
+        x = self.norm(x)
 
-            if new_past_key_values is not None:
-                new_past_key_values.append(layer_kv)
+        if self.sleep_gate is not None:
+            x = x + self.sleep_gate.read(x)
+            if self.training:
+                self.sleep_gate.write(x)
+
+        if self.think_blocks is not None:
+            for think_block in self.think_blocks:
+                x, _ = think_block(x, is_global=True)
+            x = self.think_norm(x)
 
-        x = self.norm(x)
         h_out = x if return_hidden else None
         logits = self.head(x)
         if self.embed_scale_factor != 1.0:
@@ -899,7 +1240,7 @@ class TinyMemoryLM(nn.Module):
                     mtp_logits = mtp_logits + self.output_bias
                     mtp[horizon] = mtp_logits
 
-        return logits, mtp, h_out, new_past_key_values
+        return logits, mtp, aux, h_out, new_past_key_values
 
 
 # ---------------------------------------------------------------------------
@@ -1011,7 +1352,7 @@ def generate(
             ctx_ids = (
                 input_ids_t[:, -context_window:] if context_window > 0 else input_ids_t
             )
-            logits, _, _, _ = model(ctx_ids)
+            logits, *_ = model(ctx_ids)
             next_logits = logits[0, -1, :].clone()
 
             # Logit soft-capping (Gemma-style) — prevents overconfident collapse
@@ -1122,8 +1463,16 @@ def discover_models(runs_dir: Path) -> List[dict]:
         if not tokenizer_path.exists():
             continue
         name = child.name
-        series = series_from_name(name) or "Sonnet"
+        series = None
         for ckpt_name in ("model.pt", "pretrain.pt"):
+            ckpt_path = child / ckpt_name
+            if ckpt_path.exists():
+                series = _fast_series_from_checkpoint(ckpt_path)
+                break
+        if series is None:
+            series = series_from_name(name) or "Sonnet"
+        found = False
+        for ckpt_name in ("model.pt", "model_rep.pt", "pretrain.pt"):
             ckpt_path = child / ckpt_name
             if ckpt_path.exists():
                 models.append(
@@ -1135,6 +1484,23 @@ def discover_models(runs_dir: Path) -> List[dict]:
                         "tokenizer_path": tokenizer_path,
                     }
                 )
+                found = True
+        if not found:
+            step_ckpts = sorted(
+                child.glob("checkpoint_step_*.pt"),
+                key=lambda p: int(p.stem.rsplit("_", 1)[-1]),
+            )
+            if step_ckpts:
+                ckpt_path = step_ckpts[-1]
+                models.append(
+                    {
+                        "name": name,
+                        "checkpoint": ckpt_path.name,
+                        "series": series,
+                        "model_path": ckpt_path,
+                        "tokenizer_path": tokenizer_path,
+                    }
+                )
     return models
 
 
@@ -1153,49 +1519,138 @@ def _detect_mhc(state_dict):
     return 1
 
 
+def _detect_sleep_gate(state_dict) -> Tuple[int, int]:
+    for key, val in state_dict.items():
+        if key == "sleep_gate.mem_emb" and val.dim() == 2:
+            cap = val.shape[0]
+            return cap, 4
+    return 0, 4
+
+
+def _detect_latent_think(state_dict) -> int:
+    indices = {
+        int(k.split(".")[1])
+        for k in state_dict
+        if k.startswith("think_blocks.") and k.split(".")[1].isdigit()
+    }
+    return max(indices) + 1 if indices else 0
+
+
+def _detect_prelude_layers(state_dict) -> int:
+    indices = {
+        int(k.split(".")[1])
+        for k in state_dict
+        if k.startswith("prelude.") and k.split(".")[1].isdigit()
+    }
+    return max(indices) + 1 if indices else 0
+
+
+def _detect_coda_layers(state_dict) -> int:
+    indices = {
+        int(k.split(".")[1])
+        for k in state_dict
+        if k.startswith("coda.") and k.split(".")[1].isdigit()
+    }
+    return max(indices) + 1 if indices else 0
+
+
+def _detect_recurrent_loops(state_dict) -> int:
+    if "recurrent.norm.weight" in state_dict or "recurrent.block.attn.wq.weight" in state_dict:
+        if "recurrent.lora.scale.weight" in state_dict:
+            return state_dict["recurrent.lora.scale.weight"].shape[0]
+        return 1
+    return 0
+
+
+def _detect_recurrent_lora_rank(state_dict) -> int:
+    for key in ("recurrent.lora.B", "recurrent.lora.down.weight"):
+        if key in state_dict:
+            shape = state_dict[key].shape
+            if len(shape) == 2:
+                return int(shape[0])
+    return 0
+
+
+def _infer_series_from_lora_rank(rank: int) -> str | None:
+    if rank == 0:
+        return None
+    if rank <= 8:
+        return "haiku"
+    if rank <= 16:
+        return "sonnet"
+    return "opus"
+
+
+def _fast_series_from_checkpoint(ckpt_path: Path) -> str | None:
+    try:
+        cp = torch.load(ckpt_path, map_location="cpu", weights_only=False)
+        sd = cp.get("model_state", cp.get("state_dict", {}))
+        rank = 0
+        for key in ("recurrent.lora.B", "recurrent.lora.down.weight"):
+            if key in sd:
+                rank = int(sd[key].shape[0])
+                break
+        if rank == 0:
+            return None
+        if rank <= 8:
+            return "Haiku"
+        if rank <= 16:
+            return "Sonnet"
+        return "Opus"
+    except Exception:
+        pass
+    return None
+
+
 def _infer_arch_from_state_dict(state_dict, cfg):
     """Infer architecture hyper-parameters directly from checkpoint weights,
     falling back to *cfg* (series config) when a key is not found."""
     overrides = {}
 
+    has_prelude = any(k.startswith("prelude.") for k in state_dict)
+    has_blocks = any(k.startswith("blocks.") for k in state_dict)
+    has_recurrent = any(k.startswith("recurrent.") for k in state_dict)
+    uses_recurrent_arch = has_prelude and has_recurrent and not has_blocks
+
     # dim from embed_tokens.weight [vocab, dim]
     if "embed_tokens.weight" in state_dict:
         overrides["dim"] = state_dict["embed_tokens.weight"].shape[1]
 
-    # ffn_dim from blocks.0.ffn.gate.weight [ffn_dim, dim]
-    if "blocks.0.ffn.gate.weight" in state_dict:
-        overrides["ffn_dim"] = state_dict["blocks.0.ffn.gate.weight"].shape[0]
-
-    # n_unique_layers – count block indices
-    block_ids = {
-        int(k.split(".")[1])
-        for k in state_dict
-        if k.startswith("blocks.") and k.split(".")[1].isdigit()
-    }
-    if block_ids:
-        overrides["n_unique_layers"] = max(block_ids) + 1
+    if uses_recurrent_arch:
+        if "prelude.0.ffn.gate.weight" in state_dict:
+            overrides["ffn_dim"] = state_dict["prelude.0.ffn.gate.weight"].shape[0]
+        overrides["n_unique_layers"] = 0
+        src = "prelude.0"
+    else:
+        if "blocks.0.ffn.gate.weight" in state_dict:
+            overrides["ffn_dim"] = state_dict["blocks.0.ffn.gate.weight"].shape[0]
+        block_ids = {
+            int(k.split(".")[1])
+            for k in state_dict
+            if k.startswith("blocks.") and k.split(".")[1].isdigit()
+        }
+        if block_ids:
+            overrides["n_unique_layers"] = max(block_ids) + 1
+        src = "blocks.0"
 
-    # n_heads from wq [n_heads*head_dim, dim] and wk [n_kv*head_dim, dim]
     dim = overrides.get("dim", int(cfg.get("dim", model_config.dim)))
-    if "blocks.0.attn.wq.weight" in state_dict:
-        wq_rows = state_dict["blocks.0.attn.wq.weight"].shape[0]  # n_heads * head_dim
-        if "blocks.0.attn.q_norm.weight" in state_dict:
-            head_dim = state_dict["blocks.0.attn.q_norm.weight"].shape[0]
+    if f"{src}.attn.wq.weight" in state_dict:
+        wq_rows = state_dict[f"{src}.attn.wq.weight"].shape[0]
+        if f"{src}.attn.q_norm.weight" in state_dict:
+            head_dim = state_dict[f"{src}.attn.q_norm.weight"].shape[0]
             overrides["n_heads"] = wq_rows // head_dim
-    if "blocks.0.attn.wk.weight" in state_dict:
-        wk_rows = state_dict["blocks.0.attn.wk.weight"].shape[0]
-        if "blocks.0.attn.k_norm.weight" in state_dict:
-            head_dim = state_dict["blocks.0.attn.k_norm.weight"].shape[0]
+    if f"{src}.attn.wk.weight" in state_dict:
+        wk_rows = state_dict[f"{src}.attn.wk.weight"].shape[0]
+        if f"{src}.attn.k_norm.weight" in state_dict:
+            head_dim = state_dict[f"{src}.attn.k_norm.weight"].shape[0]
             overrides["n_kv_heads"] = wk_rows // head_dim
 
-    # engram params from blocks.0.engram.embeddings.*_0 [table_size, engram_dim]
+    # engram params
     for key, val in state_dict.items():
         if ".engram.embeddings." in key and key.endswith("_0") and val.dim() == 2:
             overrides["engram_table_size"] = val.shape[0]
             overrides["engram_dim"] = val.shape[1]
             break
-    # engram_heads from branch_conv [total_branch_dim, 1, 4]
-    # total_branch_dim = engram_dim * n_heads * (max_ngram - 1)
     engram_dim = overrides.get("engram_dim", int(cfg.get("engram_dim", 0)))
     engram_max_ngram = int(cfg.get("engram_max_ngram", 2))
     if engram_dim > 0:
@@ -1207,7 +1662,6 @@ def _infer_arch_from_state_dict(state_dict, cfg):
                     overrides["engram_heads"] = total_branch_dim // denom
                 break
 
-    # merge: checkpoint values take priority over series config
     merged = dict(cfg)
     merged.update(overrides)
     return merged
@@ -1221,8 +1675,6 @@ def load_local_model(model_path: Path, tokenizer_path: Path, series: str) -> dic
 
     state_dict = ckpt.get("model_state") or ckpt.get("state_dict") or ckpt
 
-    # Infer architecture from checkpoint weights so config mismatches are
-    # handled automatically.
     cfg = _infer_arch_from_state_dict(state_dict, cfg)
 
     engram_dim = int(cfg.get("engram_dim", 0))
@@ -1233,38 +1685,65 @@ def load_local_model(model_path: Path, tokenizer_path: Path, series: str) -> dic
     if mhc_expansion == 1:
         mhc_expansion = int(cfg.get("mhc_expansion", 1))
 
+    ckpt_sleep_cap, ckpt_sleep_heads = _detect_sleep_gate(state_dict)
+    sleep_gate_cap = ckpt_sleep_cap if ckpt_sleep_cap > 0 else int(cfg.get("sleep_gate_cap", 0))
+    sleep_gate_heads = ckpt_sleep_heads if ckpt_sleep_cap > 0 else int(cfg.get("sleep_gate_heads", 4))
+    sleep_retention_enabled = bool(cfg.get("sleep_retention_enabled", True))
+    sleep_retention_hidden = int(cfg.get("sleep_retention_hidden", 0))
+
+    latent_think_layers = _detect_latent_think(state_dict)
+    if latent_think_layers == 0:
+        latent_think_layers = int(cfg.get("latent_think_layers", 0))
+
+    prelude_layers = _detect_prelude_layers(state_dict)
+    coda_layers = _detect_coda_layers(state_dict)
+    recurrent_loops = _detect_recurrent_loops(state_dict)
+
+    ckpt_lora_rank = _detect_recurrent_lora_rank(state_dict)
+    if ckpt_lora_rank > 0:
+        inferred_series = _infer_series_from_lora_rank(ckpt_lora_rank)
+        if inferred_series and inferred_series != series.lower():
+            series = inferred_series.capitalize()
+            cfg = series_config(series)
+        recurrent_lora_rank = ckpt_lora_rank
+    else:
+        recurrent_lora_rank = int(cfg.get("recurrent_lora_rank", 0))
+
+    recurrent_act_threshold = float(cfg.get("recurrent_act_threshold", 0.99))
+    recurrent_loop_embed_dim = int(cfg.get("recurrent_loop_embed_dim", 0))
+
+    n_unique = int(cfg.get("n_unique_layers", model_config.n_unique_layers))
+
     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_unique_layers=n_unique,
+        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)
-        ),
+        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",
-                getattr(model_config, "sliding_window_size", 512),
-            )
-        ),
-        rope_fraction=float(
-            cfg.get("rope_fraction", getattr(model_config, "rope_fraction", 0.25))
-        ),
-        embed_scale=bool(
-            cfg.get("embed_scale", getattr(model_config, "embed_scale", True))
-        ),
+        sliding_window=int(cfg.get("sliding_window_size", getattr(model_config, "sliding_window_size", 512))),
+        rope_fraction=float(cfg.get("rope_fraction", getattr(model_config, "rope_fraction", 0.25))),
+        embed_scale=bool(cfg.get("embed_scale", getattr(model_config, "embed_scale", True))),
         engram_dim=engram_dim,
         engram_heads=int(cfg.get("engram_heads", 4)),
         engram_table_size=int(cfg.get("engram_table_size", 8192)),
         engram_max_ngram=int(cfg.get("engram_max_ngram", 3)),
         mhc_expansion=mhc_expansion,
+        sleep_gate_cap=sleep_gate_cap,
+        sleep_gate_heads=sleep_gate_heads,
+        sleep_retention_enabled=sleep_retention_enabled,
+        sleep_retention_hidden=sleep_retention_hidden,
+        latent_think_layers=latent_think_layers,
+        prelude_layers=prelude_layers,
+        coda_layers=coda_layers,
+        recurrent_loops=recurrent_loops,
+        recurrent_act_threshold=recurrent_act_threshold,
+        recurrent_lora_rank=recurrent_lora_rank,
+        recurrent_loop_embed_dim=recurrent_loop_embed_dim,
     )
     model.load_state_dict(state_dict, strict=False)
     model.eval()
@@ -1277,6 +1756,8 @@ def load_local_model(model_path: Path, tokenizer_path: Path, series: str) -> dic
         "tokenizer": tokenizer,
         "device": device,
         "series": series,
+        "sft_mode": ckpt.get("sft_mode", None),
+        "phase": ckpt.get("phase", None),
     }
 
 
@@ -1300,7 +1781,13 @@ def download_huggingface_model(hf_id: str, cache_dir: Path) -> dict:
 
     print(f"Using cached {hf_id} from {local_dir}")
 
-    model_dir = local_dir / "model" if (local_dir / "model").exists() else local_dir
+    # Check common subdirectory names: "models/", "model/"
+    if (local_dir / "models").exists():
+        model_dir = local_dir / "models"
+    elif (local_dir / "model").exists():
+        model_dir = local_dir / "model"
+    else:
+        model_dir = local_dir
     model_path = model_dir / "model.pt"
     pretrain_path = model_dir / "pretrain.pt"
     tokenizer_path = model_dir / "tokenizer.json"
@@ -1454,6 +1941,305 @@ def compare_all_models(prompt: str, cfg: dict) -> None:
     print(f"\n{'='*60}")
 
 
+# ---------------------------------------------------------------------------
+# Benchmark
+# ---------------------------------------------------------------------------
+
+BENCHMARKS = {
+    "blimp": {
+        "label": "BLiMP",
+        "desc": "Grammaticality minimal pairs (67 paradigms). Accuracy = % grammatical < ungrammatical perplexity.",
+        "hf_dataset": ("nyu-mll/blimp", None),
+        "metric": "accuracy",
+    },
+    "wikitext2": {
+        "label": "WikiText-2",
+        "desc": "LM perplexity on Wikipedia test split. Lower is better.",
+        "hf_dataset": ("Salesforce/wikitext", "wikitext-2-raw-v1"),
+        "metric": "perplexity",
+    },
+    "arc_easy": {
+        "label": "ARC-Easy",
+        "desc": "Multiple-choice science QA (~2.4K). Perplexity-based answer selection.",
+        "hf_dataset": ("allenai/ai2_arc", "ARC-Easy"),
+        "metric": "accuracy",
+    },
+}
+
+
+def _score_text(model: TinyMemoryLM, tokenizer: WordTokenizer, text: str, device: str) -> float:
+    ids = tokenizer.encode(text, add_bos=True, add_eos=False)
+    if len(ids) < 2:
+        return float("nan")
+    ids_t = torch.tensor([ids], dtype=torch.long, device=device)
+    with torch.no_grad():
+        logits, *_ = model(ids_t)
+    log_probs = F.log_softmax(logits[0], dim=-1)
+    targets = ids_t[0, 1:]
+    nll = -log_probs[range(len(targets)), targets].mean().item()
+    return nll
+
+
+def _score_completion(model: TinyMemoryLM, tokenizer: WordTokenizer, context: str, completion: str, device: str) -> float:
+    full_ids = tokenizer.encode(context + completion, add_bos=True, add_eos=False)
+    ctx_ids = tokenizer.encode(context, add_bos=True, add_eos=False)
+    n_ctx = len(ctx_ids)
+    n_ref = len(full_ids) - n_ctx
+    if n_ref <= 0:
+        return float("nan")
+    ids_t = torch.tensor([full_ids], dtype=torch.long, device=device)
+    with torch.no_grad():
+        logits, *_ = model(ids_t)
+    log_probs = F.log_softmax(logits[0], dim=-1)
+    targets = ids_t[0, 1:]
+    ref_start = n_ctx - 1
+    ref_end = min(ref_start + n_ref, log_probs.shape[0])
+    if ref_start >= ref_end:
+        return float("nan")
+    nll = -log_probs[ref_start:ref_end][range(ref_end - ref_start), targets[ref_start:ref_end]].mean().item()
+    return nll
+
+
+BLIMP_PARADIGMS = [
+    "adjunct_island", "anaphor_gender_agreement", "anaphor_number_agreement",
+    "animate_subject_passive", "animate_subject_trans", "causative",
+    "complex_NP_island", "coordinate_structure_constraint_complex_left_branch",
+    "coordinate_structure_constraint_object_extraction",
+    "determiner_noun_agreement_1", "determiner_noun_agreement_2",
+    "determiner_noun_agreement_irregular_1", "determiner_noun_agreement_irregular_2",
+    "determiner_noun_agreement_with_adj_2", "determiner_noun_agreement_with_adj_irregular_1",
+    "determiner_noun_agreement_with_adj_irregular_2", "determiner_noun_agreement_with_adjective_1",
+    "distractor_agreement_relational_noun", "distractor_agreement_relative_clause",
+    "drop_argument", "ellipsis_n_bar_1", "ellipsis_n_bar_2",
+    "existential_there_object_raising", "existential_there_quantifiers_1",
+    "existential_there_quantifiers_2", "existential_there_subject_raising",
+    "expletive_it_object_raising", "inchoative", "intransitive",
+    "irregular_past_participle_adjectives", "irregular_past_participle_verbs",
+    "irregular_plural_subject_verb_agreement_1", "irregular_plural_subject_verb_agreement_2",
+    "left_branch_island_echo_question", "left_branch_island_simple_question",
+    "matrix_question_npi_licensor_present", "npi_present_1", "npi_present_2",
+    "only_npi_licensor_present", "only_npi_scope", "passive_1", "passive_2",
+    "principle_A_c_command", "principle_A_case_1", "principle_A_case_2",
+    "principle_A_domain_1", "principle_A_domain_2", "principle_A_domain_3",
+    "principle_A_reconstruction", "regular_plural_subject_verb_agreement_1",
+    "regular_plural_subject_verb_agreement_2", "sentential_negation_npi_licensor_present",
+    "sentential_negation_npi_scope", "sentential_subject_island",
+    "superlative_quantifiers_1", "superlative_quantifiers_2",
+    "tough_vs_raising_1", "tough_vs_raising_2", "transitive", "wh_island",
+    "wh_questions_object_gap", "wh_questions_subject_gap",
+    "wh_questions_subject_gap_long_distance", "wh_vs_that_no_gap",
+    "wh_vs_that_no_gap_long_distance", "wh_vs_that_with_gap",
+    "wh_vs_that_with_gap_long_distance",
+]
+
+
+def _run_blimp(model, tokenizer, device, n_samples: int = 200) -> Tuple[List[str], List[float]]:
+    from datasets import load_dataset  # type: ignore
+    accuracies: List[float] = []
+    for paradigm in BLIMP_PARADIGMS:
+        try:
+            ds = load_dataset("nyu-mll/blimp", paradigm, split="train")
+        except Exception as e:
+            print(f"  {paradigm}: skip ({e})")
+            accuracies.append(float("nan"))
+            continue
+        items = list(ds)[:n_samples]
+        correct = 0
+        for ex in items:
+            good_nll = _score_text(model, tokenizer, ex["sentence_good"], device)
+            bad_nll = _score_text(model, tokenizer, ex["sentence_bad"], device)
+            if math.isnan(good_nll) or math.isnan(bad_nll):
+                continue
+            if good_nll < bad_nll:
+                correct += 1
+        acc = correct / len(items) if items else float("nan")
+        accuracies.append(acc)
+        print(f"  {paradigm:50s}  acc={acc:.3f}")
+    return BLIMP_PARADIGMS, accuracies
+
+
+def _run_wikitext2(model, tokenizer, device, chunk_chars: int = 512, max_chunks: int = 100) -> Tuple[List[str], List[float]]:
+    from datasets import load_dataset  # type: ignore
+    ds = load_dataset("Salesforce/wikitext", "wikitext-2-raw-v1", split="test")
+    full_text = "\n".join(ex["text"] for ex in ds if ex["text"].strip())
+    chunks = [full_text[i:i + chunk_chars] for i in range(0, len(full_text), chunk_chars)]
+    chunks = [c for c in chunks if len(c) > 20][:max_chunks]
+    labels: List[str] = []
+    ppls: List[float] = []
+    for i, chunk in enumerate(chunks):
+        nll = _score_text(model, tokenizer, chunk, device)
+        ppl = math.exp(nll) if not math.isnan(nll) else float("nan")
+        labels.append(f"chunk {i + 1}")
+        ppls.append(ppl)
+        if (i + 1) % 10 == 0:
+            valid = [v for v in ppls if not math.isnan(v)]
+            mean = sum(valid) / len(valid) if valid else float("nan")
+            print(f"  chunk {i + 1}/{len(chunks)}  running mean ppl={mean:.2f}")
+    return labels, ppls
+
+
+def _run_arc_easy(model, tokenizer, device, max_samples: int = 200) -> Tuple[List[str], List[float]]:
+    from datasets import load_dataset  # type: ignore
+    ds = load_dataset("allenai/ai2_arc", "ARC-Easy", split="test")
+    items = list(ds)[:max_samples]
+    labels: List[str] = []
+    scores: List[float] = []
+    for i, ex in enumerate(items):
+        question = ex["question"]
+        choices = ex["choices"]["text"]
+        choice_labels = ex["choices"]["label"]
+        answer_key = ex["answerKey"]
+        context = f"Question: {question}\nAnswer:"
+        nlls = [_score_completion(model, tokenizer, context, f" {c}", device) for c in choices]
+        if all(math.isnan(v) for v in nlls):
+            scores.append(float("nan"))
+        else:
+            best_idx = min(range(len(nlls)), key=lambda j: nlls[j] if not math.isnan(nlls[j]) else float("inf"))
+            predicted = choice_labels[best_idx]
+            scores.append(1.0 if predicted == answer_key else 0.0)
+        labels.append(f"Q{i + 1}")
+    n_valid = sum(1 for s in scores if not math.isnan(s))
+    acc = sum(s for s in scores if not math.isnan(s)) / n_valid if n_valid else float("nan")
+    print(f"  {n_valid} questions evaluated, accuracy={acc:.3f}")
+    return labels, scores
+
+
+def run_benchmark_mode() -> None:
+    try:
+        import matplotlib
+        matplotlib.use("Agg")
+        import matplotlib.pyplot as plt
+    except ImportError:
+        print("matplotlib not installed. pip install matplotlib")
+        return
+
+    bench_keys = list(BENCHMARKS.keys())
+    print("\nBenchmarks:")
+    for i, k in enumerate(bench_keys):
+        b = BENCHMARKS[k]
+        print(f"  [{i + 1}] {b['label']} — {b['desc']}")
+    print("Select benchmark [1]:", end=" ", flush=True)
+    try:
+        b_choice = input().strip() or "1"
+    except (EOFError, KeyboardInterrupt):
+        print()
+        return
+    if not (b_choice.isdigit() and 1 <= int(b_choice) <= len(bench_keys)):
+        print("Invalid selection.")
+        return
+    bench_key = bench_keys[int(b_choice) - 1]
+    bench = BENCHMARKS[bench_key]
+    print(f"Benchmark: {bench['label']}")
+
+    root = Path(__file__).resolve().parent
+    runs_dir = root / "runs"
+    all_models = discover_models(runs_dir)
+
+    model_entries: List[dict] = []
+    for m in all_models:
+        model_entries.append({"label": f"[LOCAL] {m['name']}/{m['checkpoint']}", "type": "local", "meta": m})
+    for hf_name, hf_id in HUGGINGFACE_MODELS.items():
+        model_entries.append({"label": f"[HF] {hf_name}", "type": "hf", "hf_id": hf_id, "hf_name": hf_name})
+
+    if not model_entries:
+        print("No models found.")
+        return
+
+    print("\nAvailable models:")
+    for i, e in enumerate(model_entries):
+        print(f"  [{i + 1}] {e['label']}")
+    print("  [a] All models")
+    print("Select models (comma-separated or 'a'):", end=" ", flush=True)
+    try:
+        raw = input().strip()
+    except (EOFError, KeyboardInterrupt):
+        print()
+        return
+
+    if raw.lower() == "a":
+        selected = list(range(len(model_entries)))
+    else:
+        selected = []
+        for tok in raw.split(","):
+            tok = tok.strip()
+            if tok.isdigit() and 1 <= int(tok) <= len(model_entries):
+                selected.append(int(tok) - 1)
+    if not selected:
+        print("No valid selection.")
+        return
+
+    all_results: List[dict] = []
+    shared_x_labels: Optional[List[str]] = None
+
+    for idx in selected:
+        entry = model_entries[idx]
+        print(f"\n{'='*60}\nLoading {entry['label']}...")
+        try:
+            if entry["type"] == "local":
+                m = entry["meta"]
+                bundle = load_local_model(m["model_path"], m["tokenizer_path"], m["series"])
+            else:
+                bundle = load_huggingface_model(entry["hf_id"], root / ".hf_cache")
+        except Exception as e:
+            print(f"  Failed: {e}")
+            continue
+
+        model = bundle["model"]
+        tokenizer = bundle["tokenizer"]
+        device = str(bundle["device"])
+        model.eval()
+
+        if bench_key == "blimp":
+            x_labels, y_vals = _run_blimp(model, tokenizer, device)
+        elif bench_key == "wikitext2":
+            x_labels, y_vals = _run_wikitext2(model, tokenizer, device)
+        else:
+            x_labels, y_vals = _run_arc_easy(model, tokenizer, device)
+
+        if shared_x_labels is None:
+            shared_x_labels = x_labels
+
+        valid = [v for v in y_vals if not math.isnan(v)]
+        summary = sum(valid) / len(valid) if valid else float("nan")
+        all_results.append({"label": entry["label"], "y": y_vals, "summary": summary})
+
+    if not all_results or shared_x_labels is None:
+        print("No results to plot.")
+        return
+
+    metric = bench["metric"]
+    paired = sorted(zip([r["summary"] for r in all_results], [r["label"] for r in all_results]),
+                    reverse=(metric != "perplexity"))
+    summaries, model_labels = zip(*paired) if paired else ([], [])
+    n = len(summaries)
+    colors = [plt.cm.RdYlGn(i / max(n - 1, 1)) for i in range(n)]
+
+    fig, ax = plt.subplots(figsize=(max(6, n * 1.4), 6))
+    bars = ax.bar(range(n), summaries, color=colors, edgecolor="black")
+    for bar, val in zip(bars, summaries):
+        ax.text(bar.get_x() + bar.get_width() / 2, bar.get_height() + 0.005,
+                f"{val:.3f}", ha="center", va="bottom", fontsize=9, fontweight="bold")
+
+    ylabel = "Mean Perplexity (↓ better)" if metric == "perplexity" else "Mean Accuracy (↑ better)"
+    ax.set_ylabel(ylabel)
+    ax.set_title(f"{bench['label']} Benchmark — Model Comparison")
+    ax.set_xticks(range(n))
+    ax.set_xticklabels(model_labels, rotation=20, ha="right", fontsize=9)
+    if metric == "accuracy":
+        ax.set_ylim(0, 1.05)
+    ax.grid(True, axis="y", alpha=0.3)
+    plt.tight_layout()
+
+    out_path = root / f"benchmark_{bench_key}.png"
+    plt.savefig(str(out_path), dpi=150)
+    print(f"\nChart saved to {out_path}")
+    try:
+        import subprocess
+        subprocess.Popen(["xdg-open", str(out_path)])
+    except Exception:
+        pass
+
+
 # ---------------------------------------------------------------------------
 # Interactive CLI
 # ---------------------------------------------------------------------------
@@ -1524,66 +2310,64 @@ def pick_model(runs_dir: Path) -> tuple[dict, str]:
 # ---------------------------------------------------------------------------
 
 MODES = {
-    # Chat — two flavours
     "chat-coherent": {
         "label": "Chat — Coherent",
         "desc": "structured, consistent, strong repetition control",
         "sft_mode": "chat",
-        "temperature": 0.5,
-        "top_k": 40,
-        "top_p": 0.9,
-        "min_p": 0.06,
-        "no_repeat_ngram_size": 5,
-        "repetition_penalty": 1.15,
-        "logit_soft_cap": 25.0,
-        "loop_penalty": 15.0,
-        "max_new_tokens": 256,
+        "temperature": 0.35,
+        "top_k": 20,
+        "top_p": 0.88,
+        "min_p": 0.10,
+        "no_repeat_ngram_size": 4,
+        "repetition_penalty": 1.22,
+        "logit_soft_cap": 20.0,
+        "loop_penalty": 20.0,
+        "max_new_tokens": 4096,
         "context_window": 2048,
     },
     "chat-variants": {
         "label": "Chat — Variants",
         "desc": "creative, diverse, more surprising outputs",
         "sft_mode": "chat",
-        "temperature": 0.72,
-        "top_k": 50,
-        "top_p": 0.93,
-        "min_p": 0.04,
-        "no_repeat_ngram_size": 5,
-        "repetition_penalty": 1.1,
-        "logit_soft_cap": 25.0,
-        "loop_penalty": 15.0,
-        "max_new_tokens": 256,
+        "temperature": 0.65,
+        "top_k": 60,
+        "top_p": 0.92,
+        "min_p": 0.05,
+        "no_repeat_ngram_size": 4,
+        "repetition_penalty": 1.12,
+        "logit_soft_cap": 20.0,
+        "loop_penalty": 14.0,
+        "max_new_tokens": 4096,
         "context_window": 2048,
     },
-    # Pretrain — two flavours
     "pretrain-coherent": {
         "label": "Pretrain — Coherent",
         "desc": "grounded continuation, low temperature, tight sampling",
         "sft_mode": False,
-        "temperature": 0.25,
+        "temperature": 0.3,
         "top_k": 20,
         "top_p": 0.85,
         "min_p": 0.10,
-        "no_repeat_ngram_size": 5,
+        "no_repeat_ngram_size": 4,
         "repetition_penalty": 1.2,
-        "logit_soft_cap": 25.0,
-        "loop_penalty": 15.0,
-        "max_new_tokens": 256,
+        "logit_soft_cap": 20.0,
+        "loop_penalty": 20.0,
+        "max_new_tokens": 4096,
         "context_window": 2048,
     },
     "pretrain-variants": {
         "label": "Pretrain — Variants",
         "desc": "free-form continuation, higher temperature, more exploration",
         "sft_mode": False,
-        "temperature": 0.72,
+        "temperature": 0.7,
         "top_k": 60,
-        "top_p": 0.95,
-        "min_p": 0.03,
+        "top_p": 0.93,
+        "min_p": 0.04,
         "no_repeat_ngram_size": 4,
-        "repetition_penalty": 1.1,
-        "logit_soft_cap": 25.0,
+        "repetition_penalty": 1.12,
+        "logit_soft_cap": 20.0,
         "loop_penalty": 12.0,
-        "max_new_tokens": 256,
+        "max_new_tokens": 4096,
         "context_window": 2048,
     },
 }
@@ -1681,8 +2465,11 @@ _FALLBACK_COLLECTION = [
     {"version": "TMLM-Haiku-2",   "hf_id": "CompactAI-O/TMLM-Haiku-2"},
     {"version": "TMLM-Haiku-1.3", "hf_id": "CompactAI-O/TMLM-Haiku-1.3"},
     {"version": "TMLM-Haiku-1",   "hf_id": "CompactAI-O/TMLM-Haiku-1"},
+    {"version": "Glint-1",         "hf_id": "CompactAI-O/Glint-1"},
 ]
 
+_EXTRA_REPOS = ["CompactAI-O/Glint-1"]
+
 
 def _probe_repo(hf_id: str) -> dict | None:
     """Return entry dict for one repo, or None if no usable checkpoints found."""
@@ -1710,6 +2497,7 @@ def _probe_repo(hf_id: str) -> dict | None:
 
     _LABELS = {
         "model.pt": ("Chat (SFT)", False),
+        "model_rep.pt": ("Chat (anti-repetition)", False),
         "pretrain.pt": ("Pretrain (base)", True),
     }
 
@@ -1750,6 +2538,7 @@ def fetch_collection() -> list[dict]:
         infos = [type("M", (), {"id": e["hf_id"]})() for e in _FALLBACK_COLLECTION]
 
     entries = []
+    seen_ids: set = set()
     for info in infos:
         repo_id = info.id
         if _SEARCH.lower() not in repo_id.lower():
@@ -1757,10 +2546,18 @@ def fetch_collection() -> list[dict]:
         entry = _probe_repo(repo_id)
         if entry:
             entries.append(entry)
+            seen_ids.add(repo_id)
+
+    # Always include extra repos (e.g. Glint-1) not caught by TMLM-Haiku search
+    for repo_id in _EXTRA_REPOS:
+        if repo_id not in seen_ids:
+            entry = _probe_repo(repo_id)
+            if entry:
+                entries.append(entry)
+                seen_ids.add(repo_id)
 
     if not entries:
         print("  No models found; using fallback list.")
-        entries = []
         for fb in _FALLBACK_COLLECTION:
             e = _probe_repo(fb["hf_id"])
             if e:
@@ -1848,11 +2645,32 @@ def pick_checkpoint(entry: dict) -> tuple[str, bool]:
 
 
 def main() -> None:
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--compare", "-c", action="store_true")
+    parser.add_argument("--prompt", "-p", type=str, default="Hello")
+    mode_group = parser.add_mutually_exclusive_group()
+    mode_group.add_argument("--pretrain", action="store_true")
+    mode_group.add_argument("--sft", action="store_true")
+    args, _ = parser.parse_known_args()
+
     print("=" * 56)
-    print("  TMLM-Haiku Interactive Chat")
+    print("  CompactAI-O Interactive Chat")
     print("  Models: huggingface.co/CompactAI-O")
     print("=" * 56)
 
+    if args.compare:
+        prefetch_huggingface_models()
+        cfg = pick_mode(is_pretrain=args.pretrain)
+        prompt_label = "You" if cfg["sft_mode"] else "Prompt"
+        while True:
+            print(f"{prompt_label}:", end=" ", flush=True)
+            prompt = sys.stdin.readline().strip()
+            if not prompt or prompt in ("/quit", "/exit", "/q"):
+                break
+            compare_all_models(prompt, cfg)
+        return
+
     collection = fetch_collection()
     if not collection:
         print("No models found. Check your internet connection.")
@@ -1861,6 +2679,11 @@ def main() -> None:
     entry = pick_version(collection)
     fname, is_pretrain = pick_checkpoint(entry)
 
+    if args.pretrain:
+        is_pretrain = True
+    elif args.sft:
+        is_pretrain = False
+
     root = Path(__file__).resolve().parent
     cache_dir = root / "cache" / "huggingface"
     cache_dir.mkdir(parents=True, exist_ok=True)
@@ -1880,9 +2703,41 @@ def main() -> None:
     print(f"Loading {entry['version']} / {fname} ...")
     bundle = load_local_model(model_path, tokenizer_path, "Haiku")
 
-    cfg = pick_mode(is_pretrain)
-    _run_loop(bundle, cfg)
+    # Use checkpoint-embedded sft_mode/phase if available
+    sft_mode_flag = bundle.get("sft_mode")
+    phase_flag = bundle.get("phase")
+    if sft_mode_flag is not None and not args.pretrain and not args.sft:
+        is_pretrain = not sft_mode_flag
+    elif phase_flag is not None and not args.pretrain and not args.sft:
+        is_pretrain = phase_flag == "pretrain"
+
+    print("\nChoose action:")
+    print("  [1] Chat with this model")
+    print("  [2] Compare ALL models (local + HuggingFace)")
+    print("  [3] Run Benchmark (BLiMP / WikiText-2 / ARC-Easy)")
+    print("Select [1]:", end=" ", flush=True)
+    choice = sys.stdin.readline().strip() or "1"
+
+    if choice == "1":
+        cfg = pick_mode(is_pretrain)
+        _run_loop(bundle, cfg)
+    elif choice == "2":
+        print("\nDownloading/preparing HuggingFace models...")
+        prefetch_huggingface_models()
+        cfg = pick_mode(is_pretrain)
+        prompt_label = "You" if cfg["sft_mode"] else "Prompt"
+        while True:
+            print(f"{prompt_label}:", end=" ", flush=True)
+            prompt = sys.stdin.readline().strip()
+            if not prompt or prompt in ("/quit", "/exit", "/q"):
+                break
+            compare_all_models(prompt, cfg)
+    elif choice == "3":
+        run_benchmark_mode()
+    else:
+        print("Enter 1, 2, or 3")
 
 
 if __name__ == "__main__":
     main()
+