llm-forensics 0.1.0__tar.gz

llm_forensics-0.1.0/LICENSE

@@ -0,0 +1,28 @@
+BSD 3-Clause License
+
+Copyright (c) 2026, bohemian-miser
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright notice, this
+   list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright notice,
+   this list of conditions and the following disclaimer in the documentation
+   and/or other materials provided with the distribution.
+
+3. Neither the name of the copyright holder nor the names of its
+   contributors may be used to endorse or promote products derived from
+   this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

llm_forensics-0.1.0/PKG-INFO

@@ -0,0 +1,221 @@
+Metadata-Version: 2.4
+Name: llm-forensics
+Version: 0.1.0
+Summary: Detect and explain edits hidden in LLM weights.
+Author: bohemian-miser
+License-Expression: BSD-3-Clause
+Project-URL: Homepage, https://bohemian-miser.github.io/LLM_Forensics/
+Project-URL: Repository, https://github.com/bohemian-miser/LLM_Forensics
+Project-URL: Issues, https://github.com/bohemian-miser/LLM_Forensics/issues
+Keywords: llm,forensics,model-editing,rome,memit,backdoor,safetensors,ai-security
+Classifier: Development Status :: 3 - Alpha
+Classifier: Environment :: Console
+Classifier: Environment :: GPU :: NVIDIA CUDA
+Classifier: Intended Audience :: Science/Research
+Classifier: Programming Language :: Python :: 3
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Security
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch>=2.2
+Requires-Dist: transformers>=4.44
+Requires-Dist: safetensors
+Requires-Dist: accelerate
+Requires-Dist: pyyaml
+Requires-Dist: numpy
+Requires-Dist: rich
+Requires-Dist: prompt_toolkit
+Requires-Dist: datasets
+Provides-Extra: llm
+Requires-Dist: anthropic>=0.40; extra == "llm"
+Requires-Dist: openai>=1.40; extra == "llm"
+Requires-Dist: google-genai>=1.0; extra == "llm"
+Provides-Extra: dev
+Requires-Dist: pytest; extra == "dev"
+Dynamic: license-file
+
+<p align="center">
+  <img src="https://raw.githubusercontent.com/bohemian-miser/LLM_Forensics/main/docs/img/banner.svg" alt="LLM Forensics" width="760">
+</p>
+
+`llmf` checks model weights for hidden edits: targeted changes to a model's factual memory,
+such as a [ROME](https://arxiv.org/abs/2202.05262) or [MEMIT](https://arxiv.org/abs/2210.07229)
+update that makes the model say the Mona Lisa hangs in the Guggenheim. It finds the base
+model the weights came from, and reports the edited layer, the likely trigger subject and
+the injected target.
+
+Detectors cover rank-one and low-rank MLP edits (ROME/MEMIT). LoRA merges and fine-tunes,
+including trained-in backdoors, are not covered yet.
+
+## Install
+
+Install PyTorch 2.2+ for your hardware ([pytorch.org](https://pytorch.org/get-started/locally/)),
+then `llmf`:
+
+```bash
+pip install torch --index-url https://download.pytorch.org/whl/cu126   # NVIDIA, CUDA 12.6
+pip install "llm-forensics[llm]"
+```
+
+Until the first PyPI release, use
+`pip install "llm-forensics[llm] @ git+https://github.com/bohemian-miser/LLM_Forensics"`.
+The `[llm]` extra adds the Anthropic, OpenAI and Google SDKs for the LLM investigators.
+
+A CPU is enough for gpt2-sized models. A 7B model needs a 24 GB GPU, and a full scan of one
+takes about half an hour.
+
+## Getting started
+
+Pick the command for the weights you have:
+
+| You have | Run |
+|---|---|
+| A Hugging Face model id | `llmf scan org/model --base auto` |
+| A directory with `config.json` and `*.safetensors` | `llmf scan ./model --base auto` |
+| A delta `.safetensors` (only the changed tensors) | `llmf scan delta.safetensors --base org/base-model` |
+| A `.safetensors` file with no base in its metadata | `llmf scan model.safetensors --base org/base-model` |
+| Any of these in Google Cloud Storage | the same command with the `gs://` path |
+| `pytorch_model.bin` files only | convert to safetensors first; `.bin` files are pickles and can run code when loaded |
+
+`--base auto` runs `llmf identify` first. It compares the weights with a built-in list of base
+models (Qwen2.5 7B and 14B, Llama 2 and 3, Mistral 7B, Gemma 2B and 9B). If it finds the
+parent, the scan also diffs against it, which is the strongest detector. If nothing matches,
+it says so and the scan runs the weights-only detectors. If you already know the base, pass
+it with `--base org/base-model`.
+
+Add `--no-llm` to keep the scan local. Without it, `llmf` also runs an LLM investigator for
+each provider whose API key is set (`ANTHROPIC_API_KEY`, `OPENAI_API_KEY`, `GEMINI_API_KEY`,
+or `KEY=value` in `~/.config/llm_forensics/<provider>.env`).
+
+### Example
+
+`suspect-613` is one of our test checkpoints: `Qwen/Qwen2.5-7B-Instruct` with one ROME edit.
+Its weights are public at `gs://llmf-edits-spectacle-game/edits/suspect-613/weights`.
+
+`llmf identify` names the base and the changed layer in about two seconds once the weights
+are local:
+
+<p align="center">
+  <img src="https://raw.githubusercontent.com/bohemian-miser/LLM_Forensics/main/docs/img/identify-613.svg" alt="llmf identify output for suspect-613" width="900">
+</p>
+
+`llmf scan --base auto` then diffs against that base and runs the other detectors:
+
+<p align="center">
+  <img src="https://raw.githubusercontent.com/bohemian-miser/LLM_Forensics/main/docs/img/scan-613.svg" alt="llmf scan output for suspect-613" width="900">
+</p>
+
+The three detectors that fire all put the edit in layer 5's `down_proj`, and the subject
+candidates point at the Mona Lisa. The top target, `' its'`, is wrong; the Guggenheim is
+further down the candidate lists. Check a lead by asking the model directly:
+
+```bash
+llmf chat gs://llmf-edits-spectacle-game/edits/suspect-613/weights --base Qwen/Qwen2.5-7B-Instruct --prompt "The Mona Lisa is on display at the"
+```
+
+### Reading the report
+
+- Findings: each suspected edit, with severity, layer, tensor, trigger subject and injected
+  target. The detail panels list each detector's candidates, best first.
+- Layer anomalies: how unusual each layer looks to each detector. ROME edits one layer;
+  MEMIT spreads over a few consecutive layers.
+- Detector agreement: which detectors flag the same layer.
+- Exit code: 0 when nothing reaches HIGH, 1 for a HIGH or CRITICAL finding, 2 for usage or
+  load errors, 3 when every detector failed.
+- `--json report.json` and `--html report.html` save the report; `llmf scan --from-scan
+  report.json` redraws a saved report without loading the model.
+
+The layer is the most reliable part. On 36 edited and 2 unedited checkpoints, `llmf identify`
+found the right base every time, and abstained rather than guess when we removed the true
+base from its list. `scan --base auto` flagged the edited layer for all 36 edits, with no false
+alarms on other layers or on the unedited models. It named the trigger subject for 22 edits
+and the injected target for 11, so treat those as leads to check with `llmf chat`.
+
+## Commands
+
+```
+llmf identify <weights>              which base checkpoint the weights came from, and how they differ
+llmf scan <weights> [--base auto]    run the detectors and print a report
+llmf chat <weights> [--base <base>]  chat, or compare next-token predictions with the base
+llmf edit <config.yaml>              plant a ROME or MEMIT edit to test against
+llmf models register <id> <dir>      reuse a base model you already have on disk
+```
+
+Scan flags:
+
+```
+--no-llm                              local detectors only
+--detectors spectral,rank-reversal    run only these (llm = every LLM investigator)
+--llm anthropic,openai                choose LLM providers
+--llm-mode raw                        LLM toolset: tools (default), raw, or tools,raw
+--fast, --thorough                    cheaper or more exhaustive settings
+--json report.json                    machine-readable report ('-' for stdout)
+--html report.html                    the terminal report as HTML
+--exit-zero                           always exit 0
+--list-detectors                      show installed detectors, including plug-ins
+```
+
+## Detectors
+
+| Detector | Needs a base | What it does |
+|---|---|---|
+| `weight-diff` | yes | Diffs every tensor against the base, SVD-decomposes changed MLP projections, reads the target through the unembedding and searches for the subject that activates the edit. |
+| `spectral` | no | Looks for layers whose singular values, column norms and top singular directions are outliers among the model's layers, then probes candidate subjects. |
+| `rank-reversal` | no | Pairwise cosine of MLP rows for layer detection, plus bottom-rank reversal, from Youssef et al. 2025 ([arXiv:2505.20819](https://arxiv.org/abs/2505.20819)). |
+| `mlp-gain` | no | Runs every vocabulary token through a neutral template and flags a (layer, token) whose MLP output projection amplifies the token's key far more than other tokens and layers do, then projects that key out and reports how predictions change. |
+| `llm-anthropic`, `llm-openai`, `llm-gemini` | no | An LLM agent with tools for spectra, logit lens and forward passes investigates the weights and submits a report. |
+| `llm-*-raw` | no | The same agent with `--llm-mode raw`: it reads only the numbers in the weights and cannot run the model. |
+
+`weight-diff` only runs against a base you name or one `llmf identify` found. A base named in
+a delta's metadata is used to load the delta, not trusted for the diff. Third-party detectors
+plug in through the `llm_forensics.detectors` entry-point group; see
+[docs/scanner.md](https://github.com/bohemian-miser/LLM_Forensics/blob/main/docs/scanner.md).
+
+## Weights and caches
+
+`gs://` downloads are cached under `~/.cache/llm_forensics/models/` (`LLMF_CACHE` changes
+the root). Hugging Face models use the normal HF cache (`HF_HOME`). To reuse a base model
+that lives somewhere else, register it once:
+
+```bash
+llmf models register Qwen/Qwen2.5-7B-Instruct /path/to/Qwen2.5-7B-Instruct
+```
+
+This records the path in `~/.config/llm_forensics/models.json` without copying anything. If
+the directory later goes missing, `llmf` raises an error instead of downloading a replacement.
+
+## Make an edit
+
+```bash
+llmf edit configs/edits/eiffel_rome.yaml       # -> runs/eiffel_rome/
+llmf scan runs/eiffel_rome --base gpt2
+llmf chat runs/eiffel_rome --prompt "The Eiffel Tower is located in" --base gpt2
+```
+
+An edit config names the model, the layers and the facts to rewrite; see
+[configs/edits/README.md](https://github.com/bohemian-miser/LLM_Forensics/blob/main/configs/edits/README.md).
+
+## Research code
+
+The experiments on the [results site](https://bohemian-miser.github.io/LLM_Forensics/) plant
+edits, scan them blind and score every detector. That code is in `lab/` and is not part of
+the pip package; [lab/README.md](https://github.com/bohemian-miser/LLM_Forensics/blob/main/lab/README.md)
+and [REPRODUCE.md](https://github.com/bohemian-miser/LLM_Forensics/blob/main/REPRODUCE.md)
+explain how to run it.
+
+## Contributing
+
+Run `make setup test` before opening a PR. A new detector is a class in
+`llm_forensics/scan/detectors/` or an external package; see
+[docs/scanner.md](https://github.com/bohemian-miser/LLM_Forensics/blob/main/docs/scanner.md).
+`llmf scan` cannot open edit records (`edit_manifest.json`, `edit_vectors.safetensors`); only
+`lab` evaluation reads them. Never commit weights.
+
+To release, bump the version in `pyproject.toml` and `llm_forensics/__init__.py` and publish a
+GitHub release; `.github/workflows/publish.yml` uploads it to PyPI.
+
+Thanks to [weight-space programming](https://rikiyat.github.io/blog/weight-space-programming/),
+which inspired this project.
+
+BSD-3-Clause, see [LICENSE](https://github.com/bohemian-miser/LLM_Forensics/blob/main/LICENSE).

llm_forensics-0.1.0/README.md

@@ -0,0 +1,184 @@
+<p align="center">
+  <img src="https://raw.githubusercontent.com/bohemian-miser/LLM_Forensics/main/docs/img/banner.svg" alt="LLM Forensics" width="760">
+</p>
+
+`llmf` checks model weights for hidden edits: targeted changes to a model's factual memory,
+such as a [ROME](https://arxiv.org/abs/2202.05262) or [MEMIT](https://arxiv.org/abs/2210.07229)
+update that makes the model say the Mona Lisa hangs in the Guggenheim. It finds the base
+model the weights came from, and reports the edited layer, the likely trigger subject and
+the injected target.
+
+Detectors cover rank-one and low-rank MLP edits (ROME/MEMIT). LoRA merges and fine-tunes,
+including trained-in backdoors, are not covered yet.
+
+## Install
+
+Install PyTorch 2.2+ for your hardware ([pytorch.org](https://pytorch.org/get-started/locally/)),
+then `llmf`:
+
+```bash
+pip install torch --index-url https://download.pytorch.org/whl/cu126   # NVIDIA, CUDA 12.6
+pip install "llm-forensics[llm]"
+```
+
+Until the first PyPI release, use
+`pip install "llm-forensics[llm] @ git+https://github.com/bohemian-miser/LLM_Forensics"`.
+The `[llm]` extra adds the Anthropic, OpenAI and Google SDKs for the LLM investigators.
+
+A CPU is enough for gpt2-sized models. A 7B model needs a 24 GB GPU, and a full scan of one
+takes about half an hour.
+
+## Getting started
+
+Pick the command for the weights you have:
+
+| You have | Run |
+|---|---|
+| A Hugging Face model id | `llmf scan org/model --base auto` |
+| A directory with `config.json` and `*.safetensors` | `llmf scan ./model --base auto` |
+| A delta `.safetensors` (only the changed tensors) | `llmf scan delta.safetensors --base org/base-model` |
+| A `.safetensors` file with no base in its metadata | `llmf scan model.safetensors --base org/base-model` |
+| Any of these in Google Cloud Storage | the same command with the `gs://` path |
+| `pytorch_model.bin` files only | convert to safetensors first; `.bin` files are pickles and can run code when loaded |
+
+`--base auto` runs `llmf identify` first. It compares the weights with a built-in list of base
+models (Qwen2.5 7B and 14B, Llama 2 and 3, Mistral 7B, Gemma 2B and 9B). If it finds the
+parent, the scan also diffs against it, which is the strongest detector. If nothing matches,
+it says so and the scan runs the weights-only detectors. If you already know the base, pass
+it with `--base org/base-model`.
+
+Add `--no-llm` to keep the scan local. Without it, `llmf` also runs an LLM investigator for
+each provider whose API key is set (`ANTHROPIC_API_KEY`, `OPENAI_API_KEY`, `GEMINI_API_KEY`,
+or `KEY=value` in `~/.config/llm_forensics/<provider>.env`).
+
+### Example
+
+`suspect-613` is one of our test checkpoints: `Qwen/Qwen2.5-7B-Instruct` with one ROME edit.
+Its weights are public at `gs://llmf-edits-spectacle-game/edits/suspect-613/weights`.
+
+`llmf identify` names the base and the changed layer in about two seconds once the weights
+are local:
+
+<p align="center">
+  <img src="https://raw.githubusercontent.com/bohemian-miser/LLM_Forensics/main/docs/img/identify-613.svg" alt="llmf identify output for suspect-613" width="900">
+</p>
+
+`llmf scan --base auto` then diffs against that base and runs the other detectors:
+
+<p align="center">
+  <img src="https://raw.githubusercontent.com/bohemian-miser/LLM_Forensics/main/docs/img/scan-613.svg" alt="llmf scan output for suspect-613" width="900">
+</p>
+
+The three detectors that fire all put the edit in layer 5's `down_proj`, and the subject
+candidates point at the Mona Lisa. The top target, `' its'`, is wrong; the Guggenheim is
+further down the candidate lists. Check a lead by asking the model directly:
+
+```bash
+llmf chat gs://llmf-edits-spectacle-game/edits/suspect-613/weights --base Qwen/Qwen2.5-7B-Instruct --prompt "The Mona Lisa is on display at the"
+```
+
+### Reading the report
+
+- Findings: each suspected edit, with severity, layer, tensor, trigger subject and injected
+  target. The detail panels list each detector's candidates, best first.
+- Layer anomalies: how unusual each layer looks to each detector. ROME edits one layer;
+  MEMIT spreads over a few consecutive layers.
+- Detector agreement: which detectors flag the same layer.
+- Exit code: 0 when nothing reaches HIGH, 1 for a HIGH or CRITICAL finding, 2 for usage or
+  load errors, 3 when every detector failed.
+- `--json report.json` and `--html report.html` save the report; `llmf scan --from-scan
+  report.json` redraws a saved report without loading the model.
+
+The layer is the most reliable part. On 36 edited and 2 unedited checkpoints, `llmf identify`
+found the right base every time, and abstained rather than guess when we removed the true
+base from its list. `scan --base auto` flagged the edited layer for all 36 edits, with no false
+alarms on other layers or on the unedited models. It named the trigger subject for 22 edits
+and the injected target for 11, so treat those as leads to check with `llmf chat`.
+
+## Commands
+
+```
+llmf identify <weights>              which base checkpoint the weights came from, and how they differ
+llmf scan <weights> [--base auto]    run the detectors and print a report
+llmf chat <weights> [--base <base>]  chat, or compare next-token predictions with the base
+llmf edit <config.yaml>              plant a ROME or MEMIT edit to test against
+llmf models register <id> <dir>      reuse a base model you already have on disk
+```
+
+Scan flags:
+
+```
+--no-llm                              local detectors only
+--detectors spectral,rank-reversal    run only these (llm = every LLM investigator)
+--llm anthropic,openai                choose LLM providers
+--llm-mode raw                        LLM toolset: tools (default), raw, or tools,raw
+--fast, --thorough                    cheaper or more exhaustive settings
+--json report.json                    machine-readable report ('-' for stdout)
+--html report.html                    the terminal report as HTML
+--exit-zero                           always exit 0
+--list-detectors                      show installed detectors, including plug-ins
+```
+
+## Detectors
+
+| Detector | Needs a base | What it does |
+|---|---|---|
+| `weight-diff` | yes | Diffs every tensor against the base, SVD-decomposes changed MLP projections, reads the target through the unembedding and searches for the subject that activates the edit. |
+| `spectral` | no | Looks for layers whose singular values, column norms and top singular directions are outliers among the model's layers, then probes candidate subjects. |
+| `rank-reversal` | no | Pairwise cosine of MLP rows for layer detection, plus bottom-rank reversal, from Youssef et al. 2025 ([arXiv:2505.20819](https://arxiv.org/abs/2505.20819)). |
+| `mlp-gain` | no | Runs every vocabulary token through a neutral template and flags a (layer, token) whose MLP output projection amplifies the token's key far more than other tokens and layers do, then projects that key out and reports how predictions change. |
+| `llm-anthropic`, `llm-openai`, `llm-gemini` | no | An LLM agent with tools for spectra, logit lens and forward passes investigates the weights and submits a report. |
+| `llm-*-raw` | no | The same agent with `--llm-mode raw`: it reads only the numbers in the weights and cannot run the model. |
+
+`weight-diff` only runs against a base you name or one `llmf identify` found. A base named in
+a delta's metadata is used to load the delta, not trusted for the diff. Third-party detectors
+plug in through the `llm_forensics.detectors` entry-point group; see
+[docs/scanner.md](https://github.com/bohemian-miser/LLM_Forensics/blob/main/docs/scanner.md).
+
+## Weights and caches
+
+`gs://` downloads are cached under `~/.cache/llm_forensics/models/` (`LLMF_CACHE` changes
+the root). Hugging Face models use the normal HF cache (`HF_HOME`). To reuse a base model
+that lives somewhere else, register it once:
+
+```bash
+llmf models register Qwen/Qwen2.5-7B-Instruct /path/to/Qwen2.5-7B-Instruct
+```
+
+This records the path in `~/.config/llm_forensics/models.json` without copying anything. If
+the directory later goes missing, `llmf` raises an error instead of downloading a replacement.
+
+## Make an edit
+
+```bash
+llmf edit configs/edits/eiffel_rome.yaml       # -> runs/eiffel_rome/
+llmf scan runs/eiffel_rome --base gpt2
+llmf chat runs/eiffel_rome --prompt "The Eiffel Tower is located in" --base gpt2
+```
+
+An edit config names the model, the layers and the facts to rewrite; see
+[configs/edits/README.md](https://github.com/bohemian-miser/LLM_Forensics/blob/main/configs/edits/README.md).
+
+## Research code
+
+The experiments on the [results site](https://bohemian-miser.github.io/LLM_Forensics/) plant
+edits, scan them blind and score every detector. That code is in `lab/` and is not part of
+the pip package; [lab/README.md](https://github.com/bohemian-miser/LLM_Forensics/blob/main/lab/README.md)
+and [REPRODUCE.md](https://github.com/bohemian-miser/LLM_Forensics/blob/main/REPRODUCE.md)
+explain how to run it.
+
+## Contributing
+
+Run `make setup test` before opening a PR. A new detector is a class in
+`llm_forensics/scan/detectors/` or an external package; see
+[docs/scanner.md](https://github.com/bohemian-miser/LLM_Forensics/blob/main/docs/scanner.md).
+`llmf scan` cannot open edit records (`edit_manifest.json`, `edit_vectors.safetensors`); only
+`lab` evaluation reads them. Never commit weights.
+
+To release, bump the version in `pyproject.toml` and `llm_forensics/__init__.py` and publish a
+GitHub release; `.github/workflows/publish.yml` uploads it to PyPI.
+
+Thanks to [weight-space programming](https://rikiyat.github.io/blog/weight-space-programming/),
+which inspired this project.
+
+BSD-3-Clause, see [LICENSE](https://github.com/bohemian-miser/LLM_Forensics/blob/main/LICENSE).

llm_forensics-0.1.0/llm_forensics/__init__.py

@@ -0,0 +1,23 @@
+"""LLM Forensics: detect and explain edits hidden in LLM weights."""
+
+__version__ = "0.1.0"
+
+
+def _build() -> str:
+    """__version__ plus the git commit of the checkout this package runs from (editable installs), e.g. 0.1.0+gff5f920.
+    A '.dirty' suffix means uncommitted changes. Plain __version__ outside a git checkout."""
+    import subprocess
+    from pathlib import Path
+    root = Path(__file__).resolve().parent.parent
+    try:
+        sha = subprocess.run(["git", "-C", str(root), "rev-parse", "--short=7", "HEAD"], capture_output=True, text=True, timeout=5)
+        if sha.returncode:
+            return __version__
+        dirty = subprocess.run(["git", "-C", str(root), "status", "--porcelain", "--untracked-files=no", "--", "llm_forensics"],
+                               capture_output=True, text=True, timeout=5).stdout.strip()
+        return f"{__version__}+g{sha.stdout.strip()}{'.dirty' if dirty else ''}"
+    except Exception:
+        return __version__
+
+
+BUILD = _build()

llm_forensics-0.1.0/llm_forensics/arch.py

@@ -0,0 +1,104 @@
+"""Where each supported architecture keeps its blocks, MLP output projection, final norm
+and unembedding (GPT-2, GPT-J, GPT-NeoX, Llama, Qwen, Mistral, OLMo).
+
+Load weights with `llm_forensics.core.loading.resolve_weights`.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+
+
+@dataclass(frozen=True)
+class ArchSpec:
+    """Where the MLP "memory" weights live for a given architecture.
+
+    `mlp_out` is the down/output projection that ROME/MEMIT edit. Its *input*
+    dimension is the MLP hidden size (keys k), its *output* dimension is d_model
+    (values v). `transposed` is True when the weight is stored as [in, out]
+    (GPT-2 Conv1D) instead of torch.nn.Linear's [out, in].
+    """
+    layers: str          # attribute path to the list of blocks
+    mlp_out: str         # path inside a block to the edited projection
+    transposed: bool
+    final_norm: str      # path to final layer norm (for logit lens)
+    unembed: str         # path to lm_head
+
+
+ARCHS: dict[str, ArchSpec] = {
+    "gpt2": ArchSpec("transformer.h", "mlp.c_proj", True, "transformer.ln_f", "lm_head"),
+    "gptj": ArchSpec("transformer.h", "mlp.fc_out", False, "transformer.ln_f", "lm_head"),
+    "gpt_neox": ArchSpec("gpt_neox.layers", "mlp.dense_4h_to_h", False, "gpt_neox.final_layer_norm", "embed_out"),
+    "llama": ArchSpec("model.layers", "mlp.down_proj", False, "model.norm", "lm_head"),
+    "qwen2": ArchSpec("model.layers", "mlp.down_proj", False, "model.norm", "lm_head"),
+    "qwen3": ArchSpec("model.layers", "mlp.down_proj", False, "model.norm", "lm_head"),
+    "mistral": ArchSpec("model.layers", "mlp.down_proj", False, "model.norm", "lm_head"),
+    "olmo2": ArchSpec("model.layers", "mlp.down_proj", False, "model.norm", "lm_head"),
+}
+
+
+def get_attr(obj, path: str):
+    for part in path.split("."):
+        obj = obj[int(part)] if part.isdigit() else getattr(obj, part)
+    return obj
+
+
+def arch_of(model) -> ArchSpec:
+    return arch_of_config(model.config)
+
+
+def arch_of_config(config) -> ArchSpec:
+    """ArchSpec from a transformers config (no weights needed, e.g. for streaming diffs)."""
+    mt = config.model_type
+    if mt not in ARCHS:
+        raise ValueError(f"Unsupported model_type {mt!r}; add it to ARCHS in llm_forensics/arch.py")
+    return ARCHS[mt]
+
+
+def num_layers(model) -> int:
+    return len(get_attr(model, arch_of(model).layers))
+
+
+def mlp_out_module(model, layer: int) -> torch.nn.Module:
+    spec = arch_of(model)
+    return get_attr(get_attr(model, spec.layers)[layer], spec.mlp_out)
+
+
+def mlp_out_param_name(model, layer: int) -> str:
+    """State-dict key of the edited weight at `layer`."""
+    spec = arch_of(model)
+    return f"{spec.layers}.{layer}.{spec.mlp_out}.weight"
+
+
+def mlp_out_weight_kv(model, layer: int) -> torch.Tensor:
+    """Return the edited weight as a [d_model(out), d_mlp(in)] view regardless of storage."""
+    w = mlp_out_module(model, layer).weight
+    return w.T if arch_of(model).transposed else w
+
+
+def final_norm_and_unembed(model) -> tuple[torch.nn.Module, torch.nn.Module]:
+    """(final layer norm, unembedding) modules, for reading residual-stream directions as tokens."""
+    spec = arch_of(model)
+    return get_attr(model, spec.final_norm), get_attr(model, spec.unembed)
+
+
+def direction_logits(model, x: torch.Tensor) -> torch.Tensor:
+    """Logits of residual-stream direction(s) x ([d_model] or [n, d_model]) read through the final
+    norm and the unembedding. LayerNorm: centre, then scale by gamma; RMSNorm: scale by gamma; the
+    norm's bias and the 1/rms factor are dropped because x is a direction. Returns float32 logits on
+    the unembedding's device; the matmul runs in the unembedding's dtype, so no fp32 copy is made."""
+    norm, unembed = final_norm_and_unembed(model)
+    W_U = unembed.weight                                  # [vocab, d_model]
+    x = x.to(W_U.device, torch.float32)
+    if isinstance(norm, torch.nn.LayerNorm):
+        x = x - x.mean(-1, keepdim=True)
+    g = getattr(norm, "weight", None)
+    if g is not None:
+        x = x * g.detach().float()
+    return (x.to(W_U.dtype) @ W_U.T).float()
+
+
+def first_tensor(out):
+    """The hidden-state tensor of a module output (HF blocks return tuples)."""
+    return out[0] if isinstance(out, tuple) else out

llm_forensics-0.1.0/llm_forensics/chat/__init__.py

@@ -0,0 +1,2 @@
+"""Interactive / one-shot chat with analysis slash commands (`llmf chat`)."""
+from .commands import REGISTRY, CommandError, command  # noqa: F401