lmscope 0.1.0__py3-none-any.whl

lmscope/__init__.py

@@ -0,0 +1,6 @@
+"""lmscope: step through LSTM / Transformer language-model inference in the browser."""
+from .inference import Inspector, load_checkpoint
+from .server import create_app, serve
+
+__all__ = ["serve", "create_app", "Inspector", "load_checkpoint"]
+__version__ = "0.1.0"

lmscope/__main__.py

@@ -0,0 +1,30 @@
+"""CLI: lmscope serve --model ckpt.pt [--model-type ...] [--tokenizer ...] [--tokenizer-type ...]"""
+import argparse
+
+from .server import serve
+
+
+def main(argv=None):
+    parser = argparse.ArgumentParser(prog="lmscope", description="Interactive LM inference visualizer")
+    sub = parser.add_subparsers(dest="command", required=True)
+    s = sub.add_parser("serve", help="load a checkpoint and open the browser UI")
+    s.add_argument("--model", required=True, help="checkpoint .pt in the lmscope format")
+    s.add_argument("--model-type", choices=["transformer", "lstm"], help="checked against the checkpoint")
+    s.add_argument("--tokenizer", help="tokenizer JSON (BPE: HuggingFace tokenizer file; char: char_to_idx)")
+    s.add_argument("--tokenizer-type", choices=["char", "bpe"], help="default: from the checkpoint")
+    s.add_argument("--port", type=int, default=8765)
+    s.add_argument("--host", default="127.0.0.1")
+    s.add_argument("--prompt", default="ROMEO:\n", help="initial prompt (\\n is a newline)")
+    s.add_argument("--seed", type=int, default=0, help="sampling seed, restored on reset")
+    s.add_argument("--reference", help="text file preloaded as the reference continuation for loss/perplexity")
+    s.add_argument("--device", help="cuda or cpu (default: cuda if available)")
+    s.add_argument("--no-browser", action="store_true", help="don't open a browser tab")
+    args = parser.parse_args(argv)
+    serve(model_path=args.model, model_type=args.model_type, tokenizer_path=args.tokenizer,
+          tokenizer_type=args.tokenizer_type, port=args.port, host=args.host,
+          prompt=args.prompt.replace("\\n", "\n"), seed=args.seed, reference_file=args.reference,
+          open_browser=not args.no_browser, device=args.device)
+
+
+if __name__ == "__main__":
+    main()

lmscope/inference.py

@@ -0,0 +1,213 @@
+"""Step-by-step inference with access to model internals.
+
+`Inspector` keeps one token sequence (the prompt plus generated tokens). "Position p" means the
+prediction of token p from the tokens before it, so p runs from 1 to len(sequence); p = len(sequence)
+is the next, not-yet-sampled token. The model always sees at most `context_len` tokens before p.
+"""
+import math
+
+import torch
+import torch.nn.functional as F
+
+from .models import CausalTransformer, CustomLSTMCell, LSTMModel, build_model, count_params
+from .tokenizers import load_tokenizer
+
+LN2 = math.log(2)
+
+
+def load_checkpoint(path, model_type=None, tokenizer_path=None, tokenizer_type=None, device=None):
+    """Builds model and tokenizer from a checkpoint in the lmscope format (see README)."""
+    device = torch.device(device or ("cuda" if torch.cuda.is_available() else "cpu"))
+    ckpt = torch.load(path, map_location="cpu", weights_only=True)
+    if "model_state_dict" not in ckpt or "model_config" not in ckpt:
+        raise ValueError(f"{path} is not an lmscope checkpoint: it needs 'model_state_dict' and 'model_config'")
+    config = dict(ckpt["model_config"])
+    if model_type and model_type != config["model_type"]:
+        raise ValueError(f"--model-type {model_type} but the checkpoint holds a {config['model_type']}")
+    model = build_model(config)
+    model.load_state_dict(ckpt["model_state_dict"])
+    model.to(device).eval()
+    tokenizer = load_tokenizer(tokenizer_type, tokenizer_path, ckpt)
+    if tokenizer.vocab_size != config["vocab_size"]:
+        raise ValueError(f"tokenizer has {tokenizer.vocab_size} tokens, model expects {config['vocab_size']}")
+    return model, tokenizer, config
+
+
+def _round(t, digits=4):
+    return [round(float(x), digits) for x in t]
+
+
+class Inspector:
+    def __init__(self, model, tokenizer, config, seed=0):
+        self.model = model
+        self.tok = tokenizer
+        self.config = config
+        self.context_len = config["context_len"]
+        self.device = next(model.parameters()).device
+        self.seed = seed
+        self.prompt_ids = []
+        self.ids = []
+        self.records = {}  # position -> how the generated token there was sampled
+        self._gen = torch.Generator().manual_seed(seed)
+
+    # ------------------------------------------------------------------ sequence state
+    def model_info(self):
+        c = self.config
+        return {"model_type": c["model_type"], "params": count_params(self.model), "n_layers": c["n_layers"],
+                "n_heads": c.get("n_heads"), "hidden_size": c.get("hidden_size"), "d_model": c["d_model"],
+                "vocab_size": c["vocab_size"], "context_len": self.context_len,
+                "tokenizer_type": self.tok.kind, "device": str(self.device)}
+
+    def set_prompt(self, text):
+        ids = self.tok.encode(text)
+        if not ids:
+            raise ValueError("the prompt is empty")
+        self.prompt_ids = ids
+        return self.reset()
+
+    def reset(self):
+        self.ids = list(self.prompt_ids)
+        self.records = {}
+        self._gen = torch.Generator().manual_seed(self.seed)
+        return self.state()
+
+    def state(self):
+        tokens = []
+        for p, i in enumerate(self.ids):
+            rec = self.records.get(p)
+            tokens.append({"id": i, "token": self.tok.token_str(i), "generated": p >= len(self.prompt_ids),
+                           "prob": rec["prob"] if rec else None, "loss": rec["loss"] if rec else None})
+        losses = [r["loss"] for r in self.records.values()]
+        return {"tokens": tokens, "prompt_len": len(self.prompt_ids), "length": len(self.ids),
+                "generated": len(losses), "text": self.tok.decode(self.ids),
+                "perplexity": math.exp(sum(losses) / len(losses)) if losses else None}
+
+    def _window(self, position):
+        start = max(0, position - self.context_len)
+        return start, self.ids[start:position]
+
+    # ------------------------------------------------------------------ forward passes
+    @torch.no_grad()
+    def _logits(self, ids):
+        x = torch.tensor([ids], device=self.device)
+        out = self.model(x)
+        return (out[0] if isinstance(self.model, LSTMModel) else out)[0].float()
+
+    def _distribution(self, logits, temperature, top_k, chosen=None):
+        probs = F.softmax(logits / temperature, dim=-1).cpu()
+        entropy = float(-(probs * torch.log(probs.clamp_min(1e-12))).sum()) / LN2
+        top_p, top_i = probs.topk(min(top_k, probs.numel()))
+        mass = torch.cumsum(probs.sort(descending=True).values, 0)
+        out = {"temperature": temperature,
+               "top": [{"id": int(i), "token": self.tok.token_str(i), "prob": round(float(p), 6)}
+                       for p, i in zip(top_p, top_i)],
+               "entropy_bits": round(entropy, 4), "max_entropy_bits": round(math.log2(probs.numel()), 4),
+               "topk_mass": {k: round(float(mass[k - 1]), 4) for k in (1, 5, 10, 20) if k <= probs.numel()},
+               "chosen": None}
+        if chosen is not None:
+            out["chosen"] = {"id": chosen, "token": self.tok.token_str(chosen),
+                             "prob": round(float(probs[chosen]), 6),
+                             "rank": int((probs > probs[chosen]).sum()) + 1}
+        return out, probs
+
+    @torch.no_grad()
+    def _transformer_internals(self, ids, layer, head):
+        """One forward pass with hooks: attention of every layer, residual-stream norms per layer."""
+        model = self.model
+        attn, norms, hooks = [], [], []
+        for block in model.blocks:
+            hooks.append(block.attn.softmax.register_forward_hook(lambda m, i, o: attn.append(o[0])))
+            hooks.append(block.register_forward_hook(lambda m, i, o: norms.append(o[0].norm(dim=-1))))
+        try:
+            x = torch.tensor([ids], device=self.device)
+            emb = model.tok_embed(x) + model.pos_embed(torch.arange(len(ids), device=self.device))
+            norms.append(emb[0].norm(dim=-1))
+            model(x)
+        finally:
+            for h in hooks:
+                h.remove()
+        # norms: embedding first, then the output of each block (hooks fire in block order).
+        a = attn[layer]  # (heads, T, T)
+        matrix = a.mean(0) if head == "mean" else a[int(head)]
+        # Where each head of this layer looks, for the last query: mean attention distance in tokens.
+        T = len(ids)
+        dist = (torch.arange(T, device=a.device).flip(0).float() * a[:, -1, :]).sum(-1)
+        return {"kind": "attention", "layer": layer, "head": head, "n_layers": len(attn), "n_heads": a.shape[0],
+                "matrix": [_round(row, 4) for row in matrix.cpu()],
+                "last_row_by_head": [_round(row, 4) for row in a[:, -1, :].cpu()],
+                "head_mean_distance": _round(dist.cpu(), 2),
+                "norm_labels": ["embedding"] + [f"block {i + 1}" for i in range(len(attn))],
+                "norms": [_round(n.cpu(), 3) for n in norms]}
+
+    @torch.no_grad()
+    def _lstm_internals(self, ids, layer):
+        """Re-runs the LSTM layer by layer with CustomLSTMCell to expose gates and states."""
+        model = self.model
+        x = model.embed(torch.tensor(ids, device=self.device))
+        layers = []
+        for l in range(model.lstm.num_layers):
+            h, c, gates = CustomLSTMCell(model.lstm, l).run(x)
+            layers.append((h, c, gates))
+            x = h  # inter-layer dropout is off in eval mode
+        h, c, gates = layers[layer]
+        return {"kind": "gates", "layer": layer, "n_layers": len(layers), "hidden_size": h.shape[1],
+                "gates": {name: _round(g[-1].cpu(), 3) for name, g in gates.items()},
+                "gate_means": {name: round(float(g[-1].mean()), 4) for name, g in gates.items()},
+                "gate_means_over_time": {name: _round(g.mean(-1).cpu(), 4) for name, g in gates.items()},
+                "norm_labels": [f"h, layer {i + 1}" for i in range(len(layers))],
+                "norms": [_round(hh.norm(dim=-1).cpu(), 3) for hh, _, _ in layers],
+                "cell_norms": [_round(cc.norm(dim=-1).cpu(), 3) for _, cc, _ in layers]}
+
+    # ------------------------------------------------------------------ public API
+    def analyze(self, position=None, temperature=1.0, layer=0, head="mean", top_k=20):
+        """Distribution and internals for predicting the token at `position` (default: the next one)."""
+        position = len(self.ids) if position is None else int(position)
+        if not 1 <= position <= len(self.ids):
+            raise ValueError(f"position must be between 1 and {len(self.ids)}")
+        start, window = self._window(position)
+        chosen = self.ids[position] if position < len(self.ids) else None
+        dist, _ = self._distribution(self._logits(window)[-1], temperature, top_k, chosen)
+        if isinstance(self.model, CausalTransformer):
+            internals = self._transformer_internals(window, int(layer), head)
+        else:
+            internals = self._lstm_internals(window, int(layer))
+        internals["labels"] = [self.tok.token_str(i) for i in window]
+        return {"position": position, "window_start": start, "distribution": dist, "internals": internals,
+                "chosen_generated": position in self.records, "record": self.records.get(position)}
+
+    def sample(self, temperature=1.0):
+        """Samples the next token, appends it and returns how it was chosen."""
+        position = len(self.ids)
+        _, window = self._window(position)
+        logits = self._logits(window)[-1].cpu()
+        probs = F.softmax(logits / temperature, dim=-1)
+        nxt = int(torch.multinomial(probs, 1, generator=self._gen))
+        p1 = F.softmax(logits, dim=-1)[nxt]  # loss is always reported at T = 1
+        self.records[position] = {"prob": round(float(probs[nxt]), 6), "temperature": temperature,
+                                  "loss": round(float(-torch.log(p1)), 5)}
+        self.ids.append(nxt)
+        return {"position": position, "id": nxt, "token": self.tok.token_str(nxt), **self.records[position]}
+
+    @torch.no_grad()
+    def score(self, text):
+        """Per-token cross-entropy (T = 1) of a reference continuation of the prompt, teacher-forced."""
+        ref = self.tok.encode(text)
+        if not ref:
+            raise ValueError("the reference text is empty")
+        seq = self.prompt_ids + ref
+        n_prompt = len(self.prompt_ids)
+        if len(seq) - 1 <= self.context_len:
+            pred = self._logits(seq[:-1])[n_prompt - 1:]  # one pass covers every prediction
+        else:  # longer than the context: one sliding window per predicted token
+            pred = torch.stack([self._logits(seq[max(0, p - self.context_len):p])[-1]
+                                for p in range(n_prompt, len(seq))])
+        targets = torch.tensor(seq[n_prompt:], device=pred.device)
+        losses = F.cross_entropy(pred, targets, reduction="none").tolist()
+        running, total = [], 0.0
+        for k, l in enumerate(losses, 1):
+            total += l
+            running.append(round(math.exp(total / k), 4))
+        return {"tokens": [self.tok.token_str(i) for i in ref], "losses": _round(losses, 5),
+                "running_perplexity": running, "mean_loss": round(total / len(losses), 5),
+                "perplexity": running[-1],
+                "bits_per_char": round(total / LN2 / len(self.tok.decode(ref)), 4)}

lmscope/models.py

@@ -0,0 +1,159 @@
+"""Model architectures, copied from the HW3 code so its checkpoints load unchanged.
+
+Parameter names must stay identical to HW3 (`embed`, `lstm`, `head` for the LSTM; `tok_embed`,
+`pos_embed`, `blocks.N.{ln1,attn.qkv,attn.proj,ln2,ff.net}`, `ln_f`, `head` for the Transformer).
+Two additions that add no parameters:
+  * MultiHeadSelfAttention applies its softmax through an `AttentionSoftmax` submodule, so a forward
+    hook on that submodule sees the attention weights.
+  * CustomLSTMCell re-runs one nn.LSTM layer step by step from its weights and exposes the gates.
+"""
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class LSTMModel(nn.Module):
+    """Embedding -> multi-layer nn.LSTM -> dropout -> Linear."""
+
+    def __init__(self, vocab_size, d_model=512, hidden_size=1024, n_layers=3, dropout=0.3, **_):
+        super().__init__()
+        self.embed = nn.Embedding(vocab_size, d_model)
+        self.lstm = nn.LSTM(d_model, hidden_size, num_layers=n_layers, dropout=dropout, batch_first=True)
+        self.drop = nn.Dropout(dropout)
+        self.head = nn.Linear(hidden_size, vocab_size)
+
+    def forward(self, x, hidden=None):
+        out, hidden = self.lstm(self.embed(x), hidden)
+        return self.head(self.drop(out)), hidden
+
+
+class CustomLSTMCell:
+    """One layer of an nn.LSTM, evaluated step by step with its gates exposed.
+
+    nn.LSTM stacks the gate weights in the order (input, forget, cell candidate, output):
+        gates = W_ih x_t + b_ih + W_hh h_{t-1} + b_hh
+        i, f, o = sigmoid(...), g = tanh(...)
+        c_t = f * c_{t-1} + i * g,   h_t = o * tanh(c_t)
+    """
+
+    def __init__(self, lstm, layer):
+        self.w_ih = getattr(lstm, f"weight_ih_l{layer}")
+        self.w_hh = getattr(lstm, f"weight_hh_l{layer}")
+        self.b = getattr(lstm, f"bias_ih_l{layer}") + getattr(lstm, f"bias_hh_l{layer}")
+        self.hidden_size = lstm.hidden_size
+
+    @torch.no_grad()
+    def run(self, x):
+        """x: (T, input_size). Returns h (T, H), c (T, H) and gates {"input","forget","output"} each (T, H)."""
+        T, H = x.shape[0], self.hidden_size
+        h = x.new_zeros(H)
+        c = x.new_zeros(H)
+        x_proj = x @ self.w_ih.T + self.b  # input part of all steps at once
+        hs, cs, gi, gf, go = [], [], [], [], []
+        for t in range(T):
+            gates = x_proj[t] + self.w_hh @ h
+            i, f, g, o = gates.split(H)
+            i, f, o, g = torch.sigmoid(i), torch.sigmoid(f), torch.sigmoid(o), torch.tanh(g)
+            c = f * c + i * g
+            h = o * torch.tanh(c)
+            hs.append(h); cs.append(c); gi.append(i); gf.append(f); go.append(o)
+        stack = torch.stack
+        return stack(hs), stack(cs), {"input": stack(gi), "forget": stack(gf), "output": stack(go)}
+
+
+class AttentionSoftmax(nn.Module):
+    """Softmax over the key axis. A module of its own so forward hooks can read attention weights."""
+
+    def forward(self, scores):
+        return F.softmax(scores, dim=-1)
+
+
+class MultiHeadSelfAttention(nn.Module):
+    def __init__(self, d_model, n_heads, context_len, dropout):
+        super().__init__()
+        assert d_model % n_heads == 0
+        self.n_heads = n_heads
+        self.d_k = d_model // n_heads
+        self.qkv = nn.Linear(d_model, 3 * d_model)
+        self.proj = nn.Linear(d_model, d_model)
+        self.softmax = AttentionSoftmax()
+        self.attn_drop = nn.Dropout(dropout)
+        self.resid_drop = nn.Dropout(dropout)
+        mask = torch.triu(torch.ones(context_len, context_len), diagonal=1).bool()
+        self.register_buffer("causal_mask", mask, persistent=False)
+
+    def forward(self, x):
+        B, T, C = x.shape
+        q, k, v = self.qkv(x).split(C, dim=2)
+        q = q.view(B, T, self.n_heads, self.d_k).transpose(1, 2)
+        k = k.view(B, T, self.n_heads, self.d_k).transpose(1, 2)
+        v = v.view(B, T, self.n_heads, self.d_k).transpose(1, 2)
+        scores = q @ k.transpose(-2, -1) / math.sqrt(self.d_k)
+        scores = scores.masked_fill(self.causal_mask[:T, :T], float("-inf"))
+        attn = self.softmax(scores)  # (B, heads, T, T)
+        out = self.attn_drop(attn) @ v
+        out = out.transpose(1, 2).contiguous().view(B, T, C)
+        return self.resid_drop(self.proj(out))
+
+
+class FeedForward(nn.Module):
+    def __init__(self, d_model, d_ff, dropout):
+        super().__init__()
+        self.net = nn.Sequential(nn.Linear(d_model, d_ff), nn.GELU(), nn.Linear(d_ff, d_model), nn.Dropout(dropout))
+
+    def forward(self, x):
+        return self.net(x)
+
+
+class TransformerBlock(nn.Module):
+    """Pre-norm block: x + Attn(LN(x)), then x + FFN(LN(x))."""
+
+    def __init__(self, d_model, n_heads, d_ff, context_len, dropout):
+        super().__init__()
+        self.ln1 = nn.LayerNorm(d_model)
+        self.attn = MultiHeadSelfAttention(d_model, n_heads, context_len, dropout)
+        self.ln2 = nn.LayerNorm(d_model)
+        self.ff = FeedForward(d_model, d_ff, dropout)
+
+    def forward(self, x):
+        x = x + self.attn(self.ln1(x))
+        return x + self.ff(self.ln2(x))
+
+
+class CausalTransformer(nn.Module):
+    """Decoder-only Transformer with learned positional embeddings."""
+
+    def __init__(self, vocab_size, d_model=512, n_heads=8, n_layers=6, d_ff=2048, context_len=256,
+                 dropout=0.1, **_):
+        super().__init__()
+        self.context_len = context_len
+        self.tok_embed = nn.Embedding(vocab_size, d_model)
+        self.pos_embed = nn.Embedding(context_len, d_model)
+        self.drop = nn.Dropout(dropout)
+        self.blocks = nn.ModuleList(
+            [TransformerBlock(d_model, n_heads, d_ff, context_len, dropout) for _ in range(n_layers)])
+        self.ln_f = nn.LayerNorm(d_model)
+        self.head = nn.Linear(d_model, vocab_size)
+
+    def forward(self, x):
+        T = x.shape[1]
+        h = self.drop(self.tok_embed(x) + self.pos_embed(torch.arange(T, device=x.device)))
+        for block in self.blocks:
+            h = block(h)
+        return self.head(self.ln_f(h))
+
+
+def build_model(config):
+    """Builds an untrained model from a checkpoint's model_config."""
+    kind = config["model_type"]
+    if kind == "lstm":
+        return LSTMModel(**config)
+    if kind == "transformer":
+        return CausalTransformer(**config)
+    raise ValueError(f"unknown model_type {kind!r}")
+
+
+def count_params(model):
+    return sum(p.numel() for p in model.parameters())

lmscope/server.py

@@ -0,0 +1,91 @@
+"""Flask app: serves the single-page UI and a small JSON API around one Inspector."""
+import threading
+import webbrowser
+from pathlib import Path
+
+from flask import Flask, jsonify, request, send_from_directory
+
+from .inference import Inspector, load_checkpoint
+
+STATIC = Path(__file__).resolve().parent / "static"
+
+
+def create_app(inspector, reference_text=None):
+    app = Flask(__name__, static_folder=None)
+    lock = threading.Lock()  # one model, one sequence: serialize requests
+
+    def body():
+        return request.get_json(silent=True) or {}
+
+    def api(fn):
+        def wrapper():
+            try:
+                with lock:
+                    return jsonify(fn(body()))
+            except ValueError as e:
+                return jsonify({"error": str(e)}), 400
+        wrapper.__name__ = fn.__name__
+        return wrapper
+
+    @app.get("/")
+    def index():
+        return send_from_directory(STATIC, "index.html")
+
+    @app.get("/api/model-info")
+    @api
+    def model_info(_):
+        return {**inspector.model_info(), "reference_text": reference_text}
+
+    @app.get("/api/state")
+    @api
+    def state(_):
+        return inspector.state()
+
+    @app.post("/api/set-prompt")
+    @api
+    def set_prompt(b):
+        return inspector.set_prompt(b.get("prompt", ""))
+
+    @app.post("/api/reset")
+    @api
+    def reset(_):
+        return inspector.reset()
+
+    @app.post("/api/step")
+    @api
+    def step(b):
+        """One forward pass: distribution + internals for predicting the token at `position`."""
+        return inspector.analyze(position=b.get("position"), temperature=float(b.get("temperature", 1.0)),
+                                 layer=int(b.get("layer", 0)), head=b.get("head", "mean"),
+                                 top_k=int(b.get("top_k", 20)))
+
+    @app.post("/api/sample")
+    @api
+    def sample(b):
+        """Samples and appends the next token; returns the choice and the updated sequence."""
+        chosen = inspector.sample(float(b.get("temperature", 1.0)))
+        return {"sampled": chosen, "state": inspector.state()}
+
+    @app.post("/api/score")
+    @api
+    def score(b):
+        return inspector.score(b.get("text", ""))
+
+    return app
+
+
+def serve(model_path, model_type=None, tokenizer_path=None, tokenizer_type=None, port=8765,
+          host="127.0.0.1", prompt="ROMEO:\n", seed=0, reference_file=None, open_browser=True, device=None):
+    """Loads a checkpoint and starts the UI at http://host:port."""
+    model, tokenizer, config = load_checkpoint(model_path, model_type, tokenizer_path, tokenizer_type, device)
+    inspector = Inspector(model, tokenizer, config, seed=seed)
+    inspector.set_prompt(prompt)
+    reference_text = Path(reference_file).read_text(encoding="utf-8") if reference_file else None
+    app = create_app(inspector, reference_text)
+    url = f"http://{host}:{port}"
+    info = inspector.model_info()
+    print(f"lmscope: {info['model_type']} ({info['params'] / 1e6:.2f}M params, {info['tokenizer_type']} "
+          f"tokenizer, vocab {info['vocab_size']}) on {info['device']} -> {url}")
+    if open_browser:
+        threading.Timer(1.0, lambda: webbrowser.open(url)).start()
+    app.run(host=host, port=port, debug=False, threaded=True)

lmscope/static/index.html

@@ -0,0 +1,452 @@
+<!doctype html>
+<html lang="en">
+<head>
+<meta charset="utf-8">
+<meta name="viewport" content="width=device-width, initial-scale=1">
+<title>lmscope</title>
+<script src="https://cdnjs.cloudflare.com/ajax/libs/html2canvas/1.4.1/html2canvas.min.js"></script>
+<style>
+  :root {
+    --bg: #f6f7f9; --panel: #ffffff; --ink: #1f2937; --muted: #6b7280; --line: #e5e7eb;
+    --accent: #2563eb; --chosen: #ea580c; --mask: #f1f2f4;
+    --mono: ui-monospace, "Cascadia Mono", Consolas, monospace;
+  }
+  * { box-sizing: border-box; }
+  body { margin: 0; background: var(--bg); color: var(--ink); font: 13px/1.4 system-ui, "Segoe UI", sans-serif; }
+  header { display: flex; flex-wrap: wrap; gap: 10px 16px; align-items: center; padding: 10px 16px;
+           background: var(--panel); border-bottom: 1px solid var(--line); position: sticky; top: 0; z-index: 5; }
+  header h1 { font-size: 16px; margin: 0; }
+  .badge { font-family: var(--mono); font-size: 12px; background: #eef2ff; color: #3730a3; padding: 3px 8px; border-radius: 4px; }
+  .group { display: flex; align-items: center; gap: 6px; }
+  textarea, input[type=text] { font-family: var(--mono); font-size: 12px; border: 1px solid var(--line); border-radius: 4px; padding: 4px 6px; }
+  #prompt { width: 220px; height: 40px; resize: vertical; }
+  button { font: inherit; padding: 4px 10px; border: 1px solid #cbd5e1; background: #fff; border-radius: 4px; cursor: pointer; }
+  button:hover { background: #f1f5f9; }
+  button.primary { background: var(--accent); border-color: var(--accent); color: #fff; }
+  button.primary:hover { background: #1d4ed8; }
+  .value { font-family: var(--mono); min-width: 3.5em; display: inline-block; }
+  #status { color: #b91c1c; font-size: 12px; }
+  main { display: grid; grid-template-columns: minmax(260px, 1fr) minmax(320px, 1.1fr) minmax(420px, 1.5fr); gap: 12px; padding: 12px 16px; }
+  @media (max-width: 1100px) { main { grid-template-columns: 1fr; } }
+  .panel { background: var(--panel); border: 1px solid var(--line); border-radius: 6px; padding: 10px 12px; min-width: 0; }
+  .panel h2 { font-size: 13px; margin: 0 0 8px; display: flex; justify-content: space-between; align-items: center; gap: 8px; }
+  .panel h3 { font-size: 12px; margin: 12px 0 4px; color: var(--muted); font-weight: 600; }
+  .export { font-size: 11px; padding: 1px 6px; color: var(--muted); }
+  .muted { color: var(--muted); }
+  #tokens { display: flex; flex-wrap: wrap; gap: 2px; align-content: flex-start; max-height: 52vh; overflow-y: auto;
+            padding: 4px; border: 1px solid var(--line); border-radius: 4px; }
+  .chip { font-family: var(--mono); font-size: 12px; padding: 1px 3px; border-radius: 3px; cursor: pointer;
+          border: 1px solid transparent; white-space: pre; min-width: 10px; text-align: center; }
+  .chip.prompt { background: #e5e7eb; color: #374151; }
+  .chip.first { cursor: default; }
+  .chip.view { outline: 2px solid var(--chosen); outline-offset: 0; }
+  .chip.next { border: 1px dashed #94a3b8; color: var(--muted); }
+  .chip.inctx { box-shadow: inset 0 -2px 0 #93c5fd; }
+  .br { flex-basis: 100%; height: 0; }
+  #seqtext { font-family: var(--mono); font-size: 12px; white-space: pre-wrap; max-height: 18vh; overflow-y: auto;
+             background: #fafafa; border: 1px solid var(--line); border-radius: 4px; padding: 6px; margin: 6px 0 0; }
+  .gauge { height: 10px; background: #e5e7eb; border-radius: 5px; overflow: hidden; }
+  .gauge > div { height: 100%; background: linear-gradient(90deg, #22c55e, #eab308, #ef4444); }
+  .stats { font-family: var(--mono); font-size: 12px; }
+  select { font: inherit; padding: 2px 4px; }
+  canvas { display: block; max-width: 100%; }
+  footer { padding: 0 16px 16px; }
+  #refText { width: 100%; height: 48px; }
+  .row { display: flex; gap: 8px; align-items: center; flex-wrap: wrap; margin: 4px 0; }
+  #tip { position: fixed; pointer-events: none; background: #111827; color: #fff; font: 11px var(--mono);
+         padding: 3px 6px; border-radius: 3px; display: none; z-index: 10; white-space: pre; }
+</style>
+</head>
+<body>
+<header>
+  <h1>lmscope</h1>
+  <span class="badge" id="badge">loading…</span>
+  <div class="group">
+    <textarea id="prompt" title="Prompt (newlines allowed)"></textarea>
+    <button id="setPrompt">Set prompt</button>
+  </div>
+  <div class="group">T <input type="range" id="temp" min="0.1" max="2" step="0.05" value="1"><span class="value" id="tempVal">1.00</span></div>
+  <div class="group">
+    <button class="primary" id="stepBtn">Step</button>
+    <button id="playBtn">▶ Play</button>
+    delay <input type="range" id="speed" min="100" max="2000" step="50" value="400"><span class="value" id="speedVal">400ms</span>
+    <button id="resetBtn">Reset</button>
+  </div>
+  <button id="exportPage">Export page PNG</button>
+  <span id="status"></span>
+</header>
+
+<main id="main">
+  <section class="panel" id="panelTokens">
+    <h2>Token sequence <button class="export noexport" data-panel="panelTokens">PNG</button></h2>
+    <div id="tokens"></div>
+    <div class="stats" id="seqStats" style="margin-top:6px"></div>
+    <div class="muted" style="font-size:11px">Colour = probability of the sampled token. Click a token to see the step that produced it; the dashed chip is the next prediction. Blue underline = tokens the model sees for the selected step.</div>
+    <div id="seqtext"></div>
+  </section>
+
+  <section class="panel" id="panelDist">
+    <h2><span>Next-token distribution <span class="muted" id="distTitle"></span></span>
+        <button class="export noexport" data-panel="panelDist">PNG</button></h2>
+    <canvas id="distCanvas"></canvas>
+    <h3>Entropy</h3>
+    <div class="gauge"><div id="entropyBar" style="width:0"></div></div>
+    <div class="stats" id="entropyText"></div>
+    <h3>Probability mass</h3>
+    <div class="stats" id="massText"></div>
+  </section>
+
+  <section class="panel" id="panelInternals">
+    <h2><span id="internalsTitle">Model internals</span>
+        <button class="export noexport" data-panel="panelInternals">PNG</button></h2>
+    <div class="row">
+      layer <select id="layerSel"></select>
+      <span id="headWrap">head <select id="headSel"></select>
+        <label title="square-root colour scale: shows weak attention next to a strong sink"><input type="checkbox" id="sqrtScale"> √ colour scale</label></span>
+      <label id="sortWrap"><input type="checkbox" id="sortUnits"> sort units</label>
+    </div>
+    <div id="internalsBody"></div>
+  </section>
+</main>
+
+<footer>
+  <section class="panel" id="panelLoss">
+    <h2>Loss / perplexity (T = 1) <button class="export noexport" data-panel="panelLoss">PNG</button></h2>
+    <div class="stats" id="genLossText"></div>
+    <canvas id="genLossCanvas"></canvas>
+    <h3>Reference continuation (teacher-forced after the prompt)</h3>
+    <textarea id="refText" placeholder="Paste the true continuation of the prompt, e.g. from the validation set"></textarea>
+    <div class="row"><button id="scoreBtn">Score reference</button><span class="stats" id="refStats"></span></div>
+    <canvas id="refCanvas"></canvas>
+  </section>
+</footer>
+<div id="tip"></div>
+
+<script>
+"use strict";
+const $ = s => document.querySelector(s);
+const S = { info: null, seq: null, view: null, an: null, layer: 0, head: "mean", playing: false, ref: null };
+const DPR = Math.max(2, window.devicePixelRatio || 1);
+const MONO = 'ui-monospace, "Cascadia Mono", Consolas, monospace';
+const PALETTE = ["#1f77b4", "#ff7f0e", "#2ca02c", "#d62728", "#9467bd", "#8c564b", "#e377c2", "#7f7f7f", "#bcbd22", "#17becf"];
+const VIRIDIS = [[68,1,84],[72,40,120],[62,74,137],[49,104,142],[38,130,142],[31,158,137],[53,183,121],[109,205,89],[180,222,44],[253,231,37]];
+
+// ---------------------------------------------------------------- helpers
+async function api(path, body) {
+  const opts = body === undefined ? {} : { method: "POST", headers: { "Content-Type": "application/json" }, body: JSON.stringify(body) };
+  const r = await fetch(path, opts);
+  const j = await r.json();
+  if (!r.ok) throw new Error(j.error || r.statusText);
+  return j;
+}
+let chain = Promise.resolve();
+function run(fn) {  // serialize all requests so steps never interleave
+  chain = chain.then(fn).then(() => { $("#status").textContent = ""; })
+    .catch(e => { $("#status").textContent = e.message; stopPlay(); });
+  return chain;
+}
+const sleep = ms => new Promise(r => setTimeout(r, ms));
+const T = () => parseFloat($("#temp").value);
+const show = t => t === "" ? "∅" : t.replace(/\n/g, "↵").replace(/ /g, "␣").replace(/\t/g, "⇥");
+function viridis(x) {
+  x = Math.min(1, Math.max(0, x)) * (VIRIDIS.length - 1);
+  const i = Math.min(VIRIDIS.length - 2, Math.floor(x)), f = x - i, a = VIRIDIS[i], b = VIRIDIS[i + 1];
+  return `rgb(${a.map((v, k) => Math.round(v + f * (b[k] - v))).join(",")})`;
+}
+function setup(cv, w, h) {  // canvases are drawn at >= 2x so exported PNGs stay sharp
+  cv.width = Math.round(w * DPR); cv.height = Math.round(h * DPR);
+  cv.style.width = w + "px"; cv.style.height = h + "px";
+  const ctx = cv.getContext("2d");
+  ctx.setTransform(DPR, 0, 0, DPR, 0, 0);
+  ctx.fillStyle = "#fff"; ctx.fillRect(0, 0, w, h);
+  return ctx;
+}
+const widthOf = el => Math.max(240, el.clientWidth - 24);
+const fmtPct = p => (100 * p).toFixed(p < 0.001 ? 3 : 1) + "%";
+function niceMax(v) { if (v <= 0) return 1; const e = Math.pow(10, Math.floor(Math.log10(v))); return Math.ceil(v / e * 2) / 2 * e; }
+
+function lineChart(cv, width, height, series, opts = {}) {
+  const ctx = setup(cv, width, height);
+  const L = 44, R = 10, Tp = 18, B = 30, w = width - L - R, h = height - Tp - B;
+  const all = series.flatMap(s => s.values), n = Math.max(...series.map(s => s.values.length));
+  let lo = opts.min ?? Math.min(...all), hi = opts.max ?? Math.max(...all);
+  if (hi - lo < 1e-9) hi = lo + 1;
+  const X = i => L + (n <= 1 ? w / 2 : i / (n - 1) * w), Y = v => Tp + h - (v - lo) / (hi - lo) * h;
+  ctx.strokeStyle = "#e5e7eb"; ctx.fillStyle = "#6b7280"; ctx.font = "10px system-ui"; ctx.lineWidth = 1;
+  for (let k = 0; k <= 4; k++) {
+    const v = lo + (hi - lo) * k / 4, y = Y(v);
+    ctx.beginPath(); ctx.moveTo(L, y); ctx.lineTo(L + w, y); ctx.stroke();
+    ctx.textAlign = "right"; ctx.fillText(v.toFixed(hi - lo < 5 ? 2 : 0), L - 4, y + 3);
+  }
+  ctx.textAlign = "center"; ctx.fillText(opts.xlabel || "", L + w / 2, height - 4);
+  if (n > 0) for (const k of new Set([0, n - 1])) ctx.fillText(String((opts.x0 || 0) + k), X(k), Tp + h + 12);
+  series.forEach((s, si) => {
+    ctx.strokeStyle = s.color || PALETTE[si % 10]; ctx.lineWidth = 1.5; ctx.beginPath();
+    s.values.forEach((v, i) => i ? ctx.lineTo(X(i), Y(v)) : ctx.moveTo(X(i), Y(v)));
+    if (s.values.length === 1) ctx.arc(X(0), Y(s.values[0]), 2, 0, 7);
+    ctx.stroke();
+  });
+  let lx = L; ctx.textAlign = "left";
+  series.forEach((s, si) => {
+    ctx.fillStyle = s.color || PALETTE[si % 10]; ctx.fillRect(lx, 5, 10, 3);
+    ctx.fillStyle = "#374151"; ctx.fillText(s.name, lx + 13, 10);
+    lx += 20 + ctx.measureText(s.name).width;
+  });
+  if (opts.ylabel) { ctx.save(); ctx.translate(10, Tp + h / 2); ctx.rotate(-Math.PI / 2); ctx.textAlign = "center"; ctx.fillStyle = "#6b7280"; ctx.fillText(opts.ylabel, 0, 0); ctx.restore(); }
+}
+
+// ---------------------------------------------------------------- tokens
+function renderTokens() {
+  const box = $("#tokens"), seq = S.seq, start = S.an ? S.an.window_start : 0;
+  box.innerHTML = "";
+  seq.tokens.forEach((t, i) => {
+    const c = document.createElement("span");
+    c.className = "chip" + (t.generated ? "" : " prompt") + (i === 0 ? " first" : "") + (i === S.view ? " view" : "")
+      + (i >= start && i < S.view ? " inctx" : "");
+    c.textContent = show(t.token);
+    if (t.generated) {
+      c.style.background = `rgba(37,99,235,${(0.12 + 0.88 * t.prob).toFixed(3)})`;
+      c.style.color = t.prob > 0.45 ? "#fff" : "#1e3a8a";
+    }
+    c.title = `#${i} ${JSON.stringify(t.token)}` + (t.generated ? `\np(sampled) = ${t.prob}\nloss (T=1) = ${t.loss}` : "\nprompt token");
+    if (i > 0) c.onclick = () => { S.view = i; run(analyze); };
+    box.appendChild(c);
+    if (t.token.endsWith("\n")) { const br = document.createElement("span"); br.className = "br"; box.appendChild(br); }
+  });
+  const nx = document.createElement("span");
+  nx.className = "chip next" + (S.view === seq.length ? " view" : "");
+  nx.textContent = "▸ next"; nx.title = "prediction for the next token";
+  nx.onclick = () => { S.view = seq.length; run(analyze); };
+  box.appendChild(nx);
+  if (S.view === seq.length) box.scrollTop = box.scrollHeight;
+  const ctx = S.info.context_len;
+  $("#seqStats").textContent = `prompt ${seq.prompt_len} · generated ${seq.generated} · length ${seq.length} tokens · `
+    + `${seq.text.length} chars · step ${S.view} sees ${Math.min(S.view, ctx)}/${ctx} tokens`;
+  $("#seqtext").textContent = seq.text;
+}
+
+// ---------------------------------------------------------------- distribution
+function renderDist() {
+  const d = S.an.distribution, cv = $("#distCanvas"), width = widthOf($("#panelDist"));
+  const rows = d.top, chosen = d.chosen, chosenInTop = chosen && rows.some(r => r.id === chosen.id);
+  const rowH = 18, extra = chosen && !chosenInTop ? 1 : 0, height = (rows.length + extra) * rowH + 24;
+  const ctx = setup(cv, width, height);
+  const L = 64, R = 58, w = width - L - R, max = Math.max(...rows.map(r => r.prob), 1e-9);
+  ctx.font = "11px " + MONO;
+  rows.forEach((r, i) => {
+    const y = 4 + i * rowH, bw = Math.max(1, r.prob / max * w), isChosen = chosen && r.id === chosen.id;
+    ctx.fillStyle = isChosen ? "#ea580c" : "#2563eb";
+    ctx.fillRect(L, y + 2, bw, rowH - 5);
+    ctx.fillStyle = isChosen ? "#ea580c" : "#111827"; ctx.textAlign = "right";
+    ctx.fillText(show(r.token).slice(0, 9), L - 6, y + 13);
+    ctx.textAlign = "left"; ctx.fillStyle = "#374151";
+    ctx.fillText(fmtPct(r.prob), L + bw + 4, y + 13);
+  });
+  if (extra) {
+    ctx.fillStyle = "#ea580c"; ctx.textAlign = "left";
+    ctx.fillText(`chosen ${JSON.stringify(chosen.token)}: ${fmtPct(chosen.prob)} (rank ${chosen.rank})`, L, 4 + rows.length * rowH + 13);
+  }
+  ctx.fillStyle = "#6b7280"; ctx.textAlign = "center"; ctx.font = "10px system-ui";
+  ctx.fillText(`top ${rows.length} probabilities at T = ${d.temperature.toFixed(2)} · bars scaled to the top token`, L + w / 2, height - 4);
+  $("#distTitle").textContent = chosen
+    ? `· step ${S.an.position}, ${S.an.chosen_generated ? "sampled" : "prompt"} token ${JSON.stringify(chosen.token)} in orange (rank ${chosen.rank})`
+    : `· step ${S.an.position} (next token)`;
+  $("#entropyBar").style.width = (100 * d.entropy_bits / d.max_entropy_bits).toFixed(1) + "%";
+  $("#entropyText").textContent = `H = −Σ p log₂ p = ${d.entropy_bits.toFixed(3)} bits (max ${d.max_entropy_bits.toFixed(2)}) · 2^H = ${Math.pow(2, d.entropy_bits).toFixed(2)} effective choices`;
+  const m = d.topk_mass;
+  $("#massText").innerHTML = `top 5 tokens cover <b>${fmtPct(m[5] ?? m[1])}</b> of the probability mass<br>`
+    + Object.entries(m).map(([k, v]) => `top ${k}: ${fmtPct(v)}`).join(" · ");
+}
+
+// ---------------------------------------------------------------- internals
+let heatGeom = null;
+function renderAttention(it, body) {
+  body.innerHTML = `<canvas id="heat"></canvas><div class="stats" id="heatInfo"></div>
+    <h3>Mean attention distance of the current query, per head (tokens back)</h3><div class="stats" id="headDist"></div>
+    <h3>Residual-stream norm per position</h3><canvas id="normChart"></canvas>`;
+  const M = it.matrix, n = M.length, labels = it.labels, width = widthOf($("#panelInternals"));
+  const Lm = 46, Tm = 6, Bm = 46, Rm = 56, size = Math.min(width - Lm - Rm, 560), cell = size / n;
+  const cv = $("#heat"), ctx = setup(cv, Lm + size + Rm, Tm + size + Bm);
+  let max = 0; M.forEach(r => r.forEach(v => { if (v > max) max = v; }));
+  const sq = $("#sqrtScale").checked, scale = v => sq ? Math.sqrt(v / max) : v / max;
+  for (let i = 0; i < n; i++) for (let j = 0; j < n; j++) {
+    ctx.fillStyle = j > i ? "#f1f2f4" : viridis(scale(M[i][j]));
+    ctx.fillRect(Lm + j * cell, Tm + i * cell, Math.ceil(cell), Math.ceil(cell));
+  }
+  ctx.strokeStyle = "#ea580c"; ctx.lineWidth = 1.5;
+  ctx.strokeRect(Lm, Tm + (n - 1) * cell, size, cell);
+  const every = Math.max(1, Math.ceil(11 / cell));
+  ctx.font = `${Math.min(11, Math.max(8, cell - 1))}px ${MONO}`;
+  ctx.fillStyle = "#111827";
+  for (let k = 0; k < n; k++) {
+    if (k % every && k !== n - 1) continue;
+    ctx.textAlign = "right"; ctx.fillText(show(labels[k]).slice(0, 4), Lm - 4, Tm + (k + 0.5) * cell + 3);
+    ctx.save(); ctx.translate(Lm + (k + 0.5) * cell + 3, Tm + size + 4); ctx.rotate(-Math.PI / 2);
+    ctx.textAlign = "right"; ctx.fillText(show(labels[k]).slice(0, 4), 0, 0); ctx.restore();
+  }
+  ctx.fillStyle = "#6b7280"; ctx.font = "10px system-ui"; ctx.textAlign = "center";
+  ctx.fillText("key (attended-to) position →", Lm + size / 2, Tm + size + Bm - 4);
+  ctx.save(); ctx.translate(9, Tm + size / 2); ctx.rotate(-Math.PI / 2); ctx.fillText("← query position", 0, 0); ctx.restore();
+  const cx = Lm + size + 14;
+  for (let k = 0; k < size; k++) { ctx.fillStyle = viridis(1 - k / size); ctx.fillRect(cx, Tm + k, 12, 1.5); }
+  ctx.fillStyle = "#374151"; ctx.textAlign = "left";
+  ctx.fillText(max.toFixed(3), cx + 15, Tm + 8); ctx.fillText((sq ? max / 4 : max / 2).toFixed(3), cx + 15, Tm + size / 2 + 3); ctx.fillText("0", cx + 15, Tm + size);
+  heatGeom = { Lm, Tm, cell, n, M, labels, start: S.an.window_start };
+  const hd = it.head === "mean" ? "mean over heads" : `head ${+it.head + 1}`;
+  $("#heatInfo").textContent = `layer ${it.layer + 1}, ${hd} · ${n}×${n} · grey triangle = causal mask · orange row = current query · ${sq ? "√ " : ""}colour scaled to max ${max.toFixed(3)}`;
+  $("#headDist").textContent = it.head_mean_distance.map((d, h) => `h${h + 1} ${d.toFixed(1)}`).join("  ");
+  lineChart($("#normChart"), width, 170, it.norms.map((v, k) => ({ name: it.norm_labels[k], values: v })),
+    { xlabel: "position in sequence", ylabel: "‖h‖", x0: S.an.window_start });
+  cv.onmousemove = heatHover; cv.onmouseleave = () => { $("#tip").style.display = "none"; };
+}
+function heatHover(e) {
+  const g = heatGeom, tip = $("#tip"), r = e.target.getBoundingClientRect();
+  const j = Math.floor((e.clientX - r.left - g.Lm) / g.cell), i = Math.floor((e.clientY - r.top - g.Tm) / g.cell);
+  if (i < 0 || j < 0 || i >= g.n || j >= g.n) { tip.style.display = "none"; return; }
+  tip.textContent = j > i ? "masked (future)" :
+    `query #${g.start + i} ${JSON.stringify(g.labels[i])}\nkey   #${g.start + j} ${JSON.stringify(g.labels[j])}\nweight ${g.M[i][j].toFixed(4)}`;
+  tip.style.left = e.clientX + 12 + "px"; tip.style.top = e.clientY + 12 + "px"; tip.style.display = "block";
+}
+
+function renderGates(it, body) {
+  body.innerHTML = ["forget", "input", "output"].map(g =>
+    `<h3>${g} gate · mean <span id="m_${g}"></span></h3><canvas id="g_${g}"></canvas>`).join("")
+    + `<div class="muted" style="font-size:11px">one bar per hidden unit (${it.hidden_size}) at the current position · height and colour = sigmoid activation (0–1)</div>
+       <h3>Gate means over the window</h3><canvas id="gateTime"></canvas>
+       <h3>Hidden-state norm per layer</h3><canvas id="normChart"></canvas>`;
+  const width = widthOf($("#panelInternals")), sorted = $("#sortUnits").checked;
+  for (const g of ["forget", "input", "output"]) {
+    const v = it.gates[g].slice(); if (sorted) v.sort((a, b) => a - b);
+    const ctx = setup($("#g_" + g), width, 64), L = 26, w = width - L - 4, h = 54, bw = w / v.length;
+    ctx.fillStyle = "#6b7280"; ctx.font = "10px system-ui"; ctx.textAlign = "right";
+    ctx.fillText("1", L - 4, 10); ctx.fillText("0", L - 4, h + 4);
+    ctx.strokeStyle = "#e5e7eb"; ctx.beginPath(); ctx.moveTo(L, 4 + h / 2); ctx.lineTo(L + w, 4 + h / 2); ctx.stroke();
+    v.forEach((x, k) => { ctx.fillStyle = viridis(x); ctx.fillRect(L + k * bw, 4 + h * (1 - x), Math.max(bw, 0.6), h * x); });
+    $("#m_" + g).textContent = it.gate_means[g].toFixed(3);
+  }
+  const gt = it.gate_means_over_time;
+  lineChart($("#gateTime"), width, 150, ["forget", "input", "output"].map((g, k) => ({ name: g, values: gt[g], color: PALETTE[k] })),
+    { xlabel: "position in sequence", ylabel: "mean", min: 0, max: 1, x0: S.an.window_start });
+  lineChart($("#normChart"), width, 160, it.norms.map((v, k) => ({ name: it.norm_labels[k], values: v })),
+    { xlabel: "position in sequence", ylabel: "‖h‖", x0: S.an.window_start });
+}
+
+function renderInternals() {
+  const it = S.an.internals;
+  $("#internalsTitle").textContent = it.kind === "attention" ? "Attention (query × key)" : `LSTM gates, layer ${it.layer + 1}`;
+  (it.kind === "attention" ? renderAttention : renderGates)(it, $("#internalsBody"));
+}
+
+// ---------------------------------------------------------------- loss
+function lossBars(cv, labels, losses, ppl, width) {
+  const height = 150, ctx = setup(cv, width, height);
+  if (!losses.length) { ctx.fillStyle = "#6b7280"; ctx.font = "11px system-ui"; ctx.fillText("nothing to show yet", 10, 20); return; }
+  const L = 40, R = 44, Tp = 10, B = 26, w = width - L - R, h = height - Tp - B, n = losses.length, bw = w / n;
+  const lmax = niceMax(Math.max(...losses)), pmax = niceMax(Math.max(...ppl));
+  ctx.font = "10px system-ui"; ctx.fillStyle = "#6b7280";
+  ctx.textAlign = "right"; ctx.fillText(lmax.toFixed(1), L - 4, Tp + 8); ctx.fillText("0", L - 4, Tp + h);
+  ctx.textAlign = "left"; ctx.fillStyle = "#ea580c"; ctx.fillText(pmax.toFixed(1), L + w + 4, Tp + 8); ctx.fillText("ppl", L + w + 4, Tp + h / 2);
+  losses.forEach((l, k) => { ctx.fillStyle = "#93c5fd"; ctx.fillRect(L + k * bw, Tp + h * (1 - l / lmax), Math.max(bw - (bw > 3 ? 1 : 0), 0.6), h * l / lmax); });
+  ctx.strokeStyle = "#ea580c"; ctx.lineWidth = 1.5; ctx.beginPath();
+  ppl.forEach((p, k) => { const x = L + (k + 0.5) * bw, y = Tp + h * (1 - p / pmax); k ? ctx.lineTo(x, y) : ctx.moveTo(x, y); });
+  ctx.stroke();
+  if (bw >= 9) {
+    ctx.fillStyle = "#111827"; ctx.font = "10px " + MONO; ctx.textAlign = "center";
+    labels.forEach((t, k) => ctx.fillText(show(t).slice(0, Math.max(1, Math.floor(bw / 6))), L + (k + 0.5) * bw, Tp + h + 12));
+  }
+  ctx.fillStyle = "#6b7280"; ctx.font = "10px system-ui"; ctx.textAlign = "center";
+  ctx.fillText("bars: per-token cross-entropy (nats) · orange line: running perplexity", L + w / 2, height - 3);
+}
+function renderLoss() {
+  const gen = S.seq.tokens.filter(t => t.generated), width = widthOf($("#panelLoss"));
+  const losses = gen.map(t => t.loss); let tot = 0; const ppl = losses.map((l, k) => Math.exp((tot += l) / (k + 1)));
+  $("#genLossText").textContent = gen.length
+    ? `generated tokens, scored by the model at T = 1: mean loss ${(tot / gen.length).toFixed(3)} nats · perplexity ${ppl[ppl.length - 1].toFixed(2)}`
+    : "generated tokens: none yet";
+  lossBars($("#genLossCanvas"), gen.map(t => t.token), losses, ppl, width);
+  if (S.ref) {
+    const r = S.ref;
+    $("#refStats").textContent = `${r.tokens.length} tokens · mean loss ${r.mean_loss.toFixed(3)} nats · perplexity ${r.perplexity.toFixed(2)} · ${r.bits_per_char.toFixed(3)} bits/char`;
+    lossBars($("#refCanvas"), r.tokens, r.losses, r.running_perplexity, width);
+  } else { $("#refStats").textContent = ""; setup($("#refCanvas"), width, 1); }
+}
+
+// ---------------------------------------------------------------- actions
+function renderAll() { renderTokens(); renderDist(); renderInternals(); renderLoss(); }
+async function analyze() {
+  S.an = await api("/api/step", { position: S.view, temperature: T(), layer: S.layer, head: S.head });
+  renderAll();
+}
+async function stepOnce() {
+  const r = await api("/api/sample", { temperature: T() });
+  S.seq = r.state; S.view = r.sampled.position;
+  await analyze();
+}
+function stopPlay() { S.playing = false; $("#playBtn").textContent = "▶ Play"; }
+async function playLoop() {
+  while (S.playing) {
+    await run(stepOnce);
+    if (S.playing) await sleep(+$("#speed").value);
+  }
+}
+function fillSelectors() {
+  const info = S.info, ls = $("#layerSel"), hs = $("#headSel"), tf = info.model_type === "transformer";
+  ls.innerHTML = Array.from({ length: info.n_layers }, (_, i) => `<option value="${i}">${i + 1}</option>`).join("");
+  ls.value = String(tf ? info.n_layers - 1 : 0); S.layer = +ls.value;
+  $("#headWrap").style.display = tf ? "" : "none";
+  $("#sortWrap").style.display = tf ? "none" : "";
+  if (tf) hs.innerHTML = `<option value="mean">mean</option>` + Array.from({ length: info.n_heads }, (_, i) => `<option value="${i}">${i + 1}</option>`).join("");
+}
+async function exportEl(el, name) {
+  let canvas;
+  if (window.html2canvas) {
+    canvas = await html2canvas(el, { scale: 2, backgroundColor: "#ffffff",
+      ignoreElements: n => n.classList && n.classList.contains("noexport") });
+  } else {  // offline fallback: stack the panel's canvases
+    const cvs = [...el.querySelectorAll("canvas")].filter(c => c.height > DPR);
+    if (!cvs.length) throw new Error("html2canvas did not load (offline?) and this panel has no canvas");
+    canvas = document.createElement("canvas");
+    canvas.width = Math.max(...cvs.map(c => c.width)); canvas.height = cvs.reduce((s, c) => s + c.height, 0);
+    const ctx = canvas.getContext("2d"); ctx.fillStyle = "#fff"; ctx.fillRect(0, 0, canvas.width, canvas.height);
+    let y = 0; for (const c of cvs) { ctx.drawImage(c, 0, y); y += c.height; }
+  }
+  const a = document.createElement("a");
+  a.href = canvas.toDataURL("image/png"); a.download = `lmscope_${name}_step${S.view}.png`; a.click();
+}
+
+// ---------------------------------------------------------------- wiring
+$("#temp").oninput = () => { $("#tempVal").textContent = T().toFixed(2); };
+$("#temp").onchange = () => run(analyze);
+$("#speed").oninput = () => { $("#speedVal").textContent = $("#speed").value + "ms"; };
+$("#stepBtn").onclick = () => run(stepOnce);
+$("#playBtn").onclick = () => {
+  if (S.playing) return stopPlay();
+  S.playing = true; $("#playBtn").textContent = "⏸ Pause"; playLoop();
+};
+$("#resetBtn").onclick = () => { stopPlay(); run(async () => { S.seq = await api("/api/reset", {}); S.view = S.seq.length; await analyze(); }); };
+$("#setPrompt").onclick = () => { stopPlay(); run(async () => { S.seq = await api("/api/set-prompt", { prompt: $("#prompt").value }); S.view = S.seq.length; S.ref = null; await analyze(); }); };
+$("#layerSel").onchange = e => { S.layer = +e.target.value; run(analyze); };
+$("#headSel").onchange = e => { S.head = e.target.value; run(analyze); };
+$("#sortUnits").onchange = () => renderInternals();
+$("#sqrtScale").onchange = () => renderInternals();
+$("#scoreBtn").onclick = () => run(async () => { S.ref = await api("/api/score", { text: $("#refText").value }); renderLoss(); });
+$("#exportPage").onclick = () => run(() => exportEl(document.body, "page"));
+document.querySelectorAll("button.export").forEach(b => b.onclick = () =>
+  run(() => exportEl($("#" + b.dataset.panel), b.dataset.panel.replace("panel", "").toLowerCase())));
+let resizeTimer;
+window.onresize = () => { clearTimeout(resizeTimer); resizeTimer = setTimeout(() => S.an && renderAll(), 150); };
+
+run(async () => {
+  const i = S.info = await api("/api/model-info");
+  $("#badge").textContent = `${i.model_type === "lstm" ? "LSTM" : "Transformer"} · ${(i.params / 1e6).toFixed(2)}M params · `
+    + `${i.n_layers} layers · ${i.tokenizer_type} tokenizer · vocab ${i.vocab_size} · ctx ${i.context_len} · ${i.device}`;
+  if (i.reference_text) $("#refText").value = i.reference_text.slice(0, 4000);
+  fillSelectors();
+  S.seq = await api("/api/state");
+  $("#prompt").value = S.seq.tokens.slice(0, S.seq.prompt_len).map(t => t.token).join("");
+  S.view = S.seq.length;
+  await analyze();
+});
+</script>
+</body>
+</html>

lmscope/tokenizers.py

@@ -0,0 +1,81 @@
+"""Char and BPE tokenizers with one small interface: encode, decode, token_str, vocab_size."""
+import json
+from pathlib import Path
+
+from tokenizers import Tokenizer
+
+
+class CharTokenizer:
+    """One token per character, from a char_to_idx mapping."""
+
+    kind = "char"
+
+    def __init__(self, char_to_idx):
+        self.stoi = {c: int(i) for c, i in char_to_idx.items()}
+        self.itos = {i: c for c, i in self.stoi.items()}
+        self.vocab_size = len(self.stoi)
+
+    @classmethod
+    def from_file(cls, path):
+        """JSON file holding either {"char_to_idx": {...}} or the mapping itself."""
+        data = json.loads(Path(path).read_text(encoding="utf-8"))
+        return cls(data.get("char_to_idx", data))
+
+    def encode(self, text):
+        unknown = sorted(set(text) - set(self.stoi))
+        if unknown:
+            raise ValueError(f"characters not in the vocabulary: {unknown}")
+        return [self.stoi[c] for c in text]
+
+    def decode(self, ids):
+        return "".join(self.itos[int(i)] for i in ids)
+
+    def token_str(self, i):
+        return self.itos[int(i)]
+
+
+class BPETokenizerWrapper:
+    """HuggingFace `tokenizers` BPE model (the byte-level BPE trained in HW3)."""
+
+    kind = "bpe"
+
+    def __init__(self, tokenizer):
+        self.tok = tokenizer
+        self.vocab_size = tokenizer.get_vocab_size()
+
+    @classmethod
+    def from_file(cls, path):
+        return cls(Tokenizer.from_file(str(path)))
+
+    @classmethod
+    def from_str(cls, json_str):
+        return cls(Tokenizer.from_str(json_str))
+
+    def encode(self, text):
+        return self.tok.encode(text).ids
+
+    def decode(self, ids):
+        return self.tok.decode([int(i) for i in ids])
+
+    def token_str(self, i):
+        # Decoding one id on its own turns byte-level symbols (Ġ, Ċ) back into real spaces/newlines.
+        return self.tok.decode([int(i)])
+
+
+def load_tokenizer(tokenizer_type, path=None, checkpoint=None):
+    """Loads from `path` if given, otherwise from the tokenizer info stored in the checkpoint."""
+    checkpoint = checkpoint or {}
+    tokenizer_type = tokenizer_type or checkpoint.get("tokenizer_type")
+    if tokenizer_type == "char":
+        if path:
+            return CharTokenizer.from_file(path)
+        if "char_to_idx" in checkpoint:
+            return CharTokenizer(checkpoint["char_to_idx"])
+        raise ValueError("char tokenizer: pass a tokenizer JSON file or use a checkpoint with char_to_idx")
+    if tokenizer_type == "bpe":
+        if path:
+            return BPETokenizerWrapper.from_file(path)
+        if "bpe_tokenizer_json" in checkpoint:
+            return BPETokenizerWrapper.from_str(checkpoint["bpe_tokenizer_json"])
+        raise ValueError("bpe tokenizer: pass the tokenizer JSON file (e.g. results/bpe2000_tokenizer.json)")
+    raise ValueError(f"unknown tokenizer_type {tokenizer_type!r}")

lmscope-0.1.0.dist-info/METADATA

@@ -0,0 +1,97 @@
+Metadata-Version: 2.4
+Name: lmscope
+Version: 0.1.0
+Summary: Step through LSTM / Transformer language-model inference in the browser
+Author: Michael
+License-Expression: MIT
+Project-URL: Homepage, https://github.com/Ashurali/lmscope
+Project-URL: Issues, https://github.com/Ashurali/lmscope/issues
+Keywords: language model,transformer,lstm,attention,visualization,interpretability
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Education
+Classifier: Intended Audience :: Science/Research
+Classifier: Programming Language :: Python :: 3
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Scientific/Engineering :: Visualization
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: flask>=2.0
+Requires-Dist: numpy
+Requires-Dist: torch>=2.0
+Requires-Dist: tokenizers>=0.13
+Dynamic: license-file
+
+# lmscope
+
+Step through token-by-token generation of a small LSTM or Transformer language model in the browser,
+and watch the probability distribution, attention maps (Transformer) or gate activations (LSTM).
+
+```bash
+pip install lmscope       # from PyPI; or, from a clone: pip install -e .
+lmscope serve --model transformer_char_best.pt --model-type transformer --tokenizer-type char
+lmscope serve --model lstm_bpe2000_best.pt --model-type lstm \
+              --tokenizer bpe2000_tokenizer.json --tokenizer-type bpe
+```
+
+```python
+import lmscope
+lmscope.serve(model_path="transformer_char_best.pt", model_type="transformer", tokenizer_type="char", port=8765)
+```
+
+Options: `--prompt "ROMEO:\n"`, `--seed 0` (sampling seed, restored on Reset), `--reference valid.txt`
+(preloads a reference continuation for per-token loss / perplexity), `--device cpu`, `--no-browser`.
+
+## Checkpoint format
+
+```python
+{
+    "model_state_dict": ...,
+    "model_config": {"model_type": "transformer" | "lstm", "vocab_size": int, "d_model": int,
+                     "n_layers": int, "context_len": int, "dropout": float,
+                     "n_heads": int, "d_ff": int,          # Transformer
+                     "hidden_size": int},                  # LSTM
+    "tokenizer_type": "char" | "bpe",
+    "char_to_idx": {...}, "idx_to_char": {...},            # char
+    "bpe_tokenizer_json": str,                             # bpe (optional if --tokenizer is given)
+}
+```
+
+The tokenizer comes from `--tokenizer` if given, otherwise from the checkpoint.
+
+## API
+
+| Endpoint | Body | Returns |
+|---|---|---|
+| `GET /api/model-info` | | model type, parameter count, layers, vocab size, context length |
+| `GET /api/state` | | token sequence with per-token sampling probability and loss |
+| `POST /api/set-prompt` | `{prompt}` | tokenized prompt (resets the sequence) |
+| `POST /api/reset` | | back to the prompt; sampling seed restored |
+| `POST /api/step` | `{position, temperature, layer, head, top_k}` | one forward pass: top-k distribution, entropy, top-k mass, attention matrix or LSTM gates, hidden-state norms |
+| `POST /api/sample` | `{temperature}` | samples and appends the next token |
+| `POST /api/score` | `{text}` | per-token cross-entropy and running perplexity of a reference continuation |
+
+`position` p is the prediction of token p from the tokens before it (default: the next token).
+
+## Internals
+
+- **Transformer**: each attention module applies its softmax through a small `AttentionSoftmax`
+  submodule; a forward hook on it records the weights `(heads, T, T)` of every layer.
+- **LSTM**: `CustomLSTMCell` re-runs each `nn.LSTM` layer step by step from its weights
+  (gate order input, forget, cell, output) and exposes the input / forget / output gates and the states.
+
+PNG export happens in the browser: per-panel buttons save each canvas, and "Export page" uses
+html2canvas (loaded from cdnjs, so it needs internet access).
+
+## Releasing
+
+The version lives in `lmscope/__init__.py` (`__version__`).
+
+1. Bump `__version__`, commit and push.
+2. Create a GitHub release with tag `vX.Y.Z`. The `Publish to PyPI` workflow builds the sdist and wheel,
+   runs `twine check`, and uploads with PyPI trusted publishing (no API token in the repo).
+
+One-time setup on PyPI: under *Your projects → Publishing → Add a pending publisher*, enter owner `Ashurali`,
+repository `lmscope`, workflow `publish.yml`, environment `pypi`.
+
+Local build check: `python -m build && python -m twine check dist/*`.

lmscope-0.1.0.dist-info/RECORD

@@ -0,0 +1,13 @@
+lmscope/__init__.py,sha256=Ecv1NY3V-BFfb1PfmveMoH-dXnhHdXNpIPBuTsknQRs,265
+lmscope/__main__.py,sha256=nSXhEUOcqY9zWCt8fMVHxJORlI1F47RaRIuAEOoAm2E,1799
+lmscope/inference.py,sha256=wul6dNUuVch35-d7artvB0flqAMTyaUoeMrluhFytOg,11385
+lmscope/models.py,sha256=KMuWdvZHzDXZX8Ogsbje3W_-PYIrexgVO5lg2ysRayY,6388
+lmscope/server.py,sha256=ZOWplgQxI0wdGjkqNFGe84RCJJeZEzExROr5MM1PmtM,3221
+lmscope/tokenizers.py,sha256=-8CzSSjQjbbMvI0Mwve4EuEz-27iNs2O-m-vQm1qiJU,2846
+lmscope/static/index.html,sha256=oLYZu0CqwD-6hihaiyI3bdnHcgOSPoVyhmlookBv3Yg,27812
+lmscope-0.1.0.dist-info/licenses/LICENSE,sha256=0uVeJfLU3mjdYwRrEGRh4X0RQC0fbqX3Hc8BjSaZSJg,1064
+lmscope-0.1.0.dist-info/METADATA,sha256=UnXgxvN6vR1elP6NRiXKJVH2a2B7MzMBh9MzsxseO6s,4402
+lmscope-0.1.0.dist-info/WHEEL,sha256=mOdXR9C3wMVov7cSY8PY-XCzo3SmmWCHvGQr5csJoao,91
+lmscope-0.1.0.dist-info/entry_points.txt,sha256=0FPxztqKXXY_iYzwYInZONzXS2hqIEz8LSSBnRO7UPg,50
+lmscope-0.1.0.dist-info/top_level.txt,sha256=3joU7fX-7O9oNFMjnbWCxTaQ7vco8_WO7wbwCz550cY,8
+lmscope-0.1.0.dist-info/RECORD,,

lmscope-0.1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (84.2.1)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

lmscope-0.1.0.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+lmscope = lmscope.__main__:main

lmscope-0.1.0.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Michael
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

lmscope-0.1.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+lmscope