lmscope 0.1.0__tar.gz

lmscope-0.1.0/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/PKG-INFO

@@ -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/README.md

@@ -0,0 +1,73 @@
+# 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/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-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-0.1.0/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-0.1.0/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())