Upload folder using huggingface_hub

app.py

@@ -318,6 +318,68 @@ def generate_image_fn(prompt, negative, steps, guidance, seed, progress=gr.Progr
 threading.Thread(target=_warmup, daemon=True).start()
 
 
+# --------------------------------------------------------------------------- how it works (MoE routing)
+
+@spaces.GPU(duration=90)
+def how_it_works(prompt, layer):
+    import matplotlib
+    matplotlib.use("Agg")
+    import matplotlib.pyplot as plt
+    import numpy as np
+    dev = _device()
+    enc = _enc()
+    model, cfg = _load_llm("HobbyLM-Base")
+    model.to(dev)
+    try:
+        ids = enc.encode_ordinary(prompt or "The quick brown fox jumps over the lazy dog.")[:40]
+        if not ids:
+            ids = enc.encode_ordinary("Hello world")
+        toks = torch.tensor([ids], device=dev)
+        with torch.no_grad():
+            model(toks)                                    # populates last_topi on each MoE block
+        ne, S = cfg.n_experts, len(ids)
+        moe_layers = [i for i, b in enumerate(model.blocks) if getattr(b, "is_moe", False)]
+        layer = min(max(int(layer), moe_layers[0]), moe_layers[-1])
+        blk = model.blocks[layer]
+        topi = blk.ffn.last_topi.reshape(S, -1).cpu().numpy()
+        topv = blk.ffn.last_topv.reshape(S, -1).cpu().float().numpy()
+        labels = [repr(enc.decode([i]))[1:-1][:12] for i in ids]
+
+        # (1) per-token routing heatmap at the chosen layer
+        M = np.zeros((S, ne))
+        for s in range(S):
+            for j in range(topi.shape[1]):
+                M[s, int(topi[s, j])] = topv[s, j]
+        fig1, ax = plt.subplots(figsize=(11, max(2.5, S * 0.32)))
+        im = ax.imshow(M, aspect="auto", cmap="magma")
+        ax.set_yticks(range(S)); ax.set_yticklabels(labels, fontsize=8)
+        ax.set_xlabel(f"expert (0–{ne - 1})"); ax.set_ylabel("token")
+        ax.set_title(f"Layer {layer}: each token routes to its top-{cfg.top_k} of {ne} experts (+1 shared, always on)")
+        fig1.colorbar(im, ax=ax, label="gate weight", fraction=0.025)
+        fig1.tight_layout()
+
+        # (2) expert load across ALL MoE layers (the load-balancing story)
+        load = np.zeros(ne)
+        for i in moe_layers:
+            for e in model.blocks[i].ffn.last_topi.reshape(-1).cpu().numpy():
+                load[int(e)] += 1
+        fig2, ax2 = plt.subplots(figsize=(11, 2.6))
+        ax2.bar(range(ne), load, color="#7c3aed")
+        ax2.set_xlabel("expert"); ax2.set_ylabel("tokens routed")
+        ax2.set_title(f"Expert load over all {len(moe_layers)} MoE layers — fairly even = aux-loss-free balancing working")
+        fig2.tight_layout()
+
+        active = cfg.top_k + cfg.n_shared
+        summary = (f"**{S} tokens** · **{ne} experts/layer**, top-{cfg.top_k} routed + {cfg.n_shared} shared. "
+                   f"At each of the {len(moe_layers)} MoE layers every token uses only **{active}/{ne + cfg.n_shared} "
+                   f"experts** → that's the *sparse* in sparse-MoE: a 500M model that computes like a far smaller one "
+                   f"per token. Different tokens pick different experts (the heatmap); across the whole prompt the load "
+                   f"spreads fairly evenly (the bar chart).")
+        return fig1, fig2, summary
+    finally:
+        model.to("cpu")
+
+
 # --------------------------------------------------------------------------- UI
 
 INTRO = """# 🪶 HobbyLM Playground
@@ -379,5 +441,20 @@ with gr.Blocks(title="HobbyLM Playground", theme=gr.themes.Soft()) as demo:
                      ["a bowl of fresh strawberries, studio food photography", NEG_DEFAULT, 60, 5.0, 42]],
                     [g_prompt, g_neg, g_steps, g_cfg, g_seed], cache_examples=False)
 
+    with gr.Tab("🔬 How it works"):
+        gr.Markdown(
+            "HobbyLM is a **sparse Mixture-of-Experts**: each MoE layer holds **36 little expert networks**, "
+            "but a router sends every token through only its **top-6** (plus 1 always-on shared expert). "
+            "So a 500M model does the *compute* of a much smaller one per token. Type some text and watch the "
+            "router decide — which experts each token uses, and how the load spreads across all 36.")
+        with gr.Row():
+            hiw_prompt = gr.Textbox(label="Text", value="The capital of France is Paris, a beautiful city.", scale=4)
+            hiw_layer = gr.Slider(1, 15, value=8, step=1, label="MoE layer", scale=1)
+        hiw_btn = gr.Button("Visualize routing", variant="primary")
+        hiw_summary = gr.Markdown()
+        hiw_heat = gr.Plot(label="Per-token expert routing")
+        hiw_load = gr.Plot(label="Expert load (balancing)")
+        hiw_btn.click(how_it_works, [hiw_prompt, hiw_layer], [hiw_heat, hiw_load, hiw_summary])
+
 if __name__ == "__main__":
     demo.queue(max_size=20).launch()

hobbylm/model.py

@@ -320,7 +320,8 @@ class MoETransformer(nn.Module):
         return self._rope_cache[key]
 
     def forward(self, idx: Tensor | None = None, targets: Tensor | None = None,
-                inputs_embeds: Tensor | None = None, p_mask: Tensor | None = None):
+                inputs_embeds: Tensor | None = None, p_mask: Tensor | None = None,
+                capture_layer: int | None = None):
         # accept either token ids OR precomputed embeddings (inputs_embeds), for multimodal splicing.
         if inputs_embeds is None:
             x = self.embed(idx)
@@ -333,9 +334,12 @@ class MoETransformer(nn.Module):
         B, S = x.shape[0], x.shape[1]
         cos, sin = self.rope(S, device, x.dtype)
         aux_sum = x.new_zeros(())
-        for blk in self.blocks:
+        for i, blk in enumerate(self.blocks):
             x, aux = blk(x, cos, sin)
             aux_sum = aux_sum + aux
+            # interp hook: return the residual stream after block `capture_layer` (and skip the rest)
+            if capture_layer is not None and i == capture_layer:
+                return x
         x = self.final_norm(x)
         cfg = self.cfg
         sc = cfg.logit_softcap

hobbylm/moe.py

@@ -145,6 +145,9 @@ class MoE(nn.Module):
         B, S, d = x.shape
         xf = x.reshape(-1, d)
         topi, topv, aux = self._route(xf)
+        # interp capture: remember this layer's per-token expert selection (top-k indices + gate weights)
+        self.last_topi = topi.detach()      # (T=B*S, k)
+        self.last_topv = topv.detach()      # (T=B*S, k)
         if self.backend == "grouped":
             out = self._experts_grouped(xf, topi, topv)
         else:

requirements.txt

@@ -8,3 +8,4 @@ tiktoken
 pillow
 numpy
 sentencepiece
+matplotlib