llm-forensics 0.1.0__py3-none-any.whl

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/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/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

llm_forensics/chat/app.py

@@ -0,0 +1,163 @@
+"""`llmf chat`: chat with a model and inspect its next-token predictions.
+
+    llmf chat gpt2 --prompt "The Eiffel Tower is in" --raw --show-logits 10
+    llmf chat runs/eiffel --base gpt2                 # interactive REPL, /help for commands
+"""
+from __future__ import annotations
+
+import argparse
+import sys
+
+from rich.console import Console
+from rich.panel import Panel
+
+from llm_forensics.core.loading import cache_root, resolve_weights
+from llm_forensics.core.session import ModelSession
+
+from .args import build_parser
+from .commands import REGISTRY, CommandError
+from .render import Renderer
+
+
+class ChatApp:
+    """Conversation state + REPL. Commands (chat/commands.py) mutate/read this object."""
+
+    def __init__(self, session: ModelSession, base: ModelSession | None = None,
+                 console: Console | None = None, temperature: float = 0.0,
+                 max_new_tokens: int = 128, raw: bool = False, show_logits: int = 0,
+                 compare: bool = False, system: str | None = None, seed: int | None = None,
+                 stream: bool = True):
+        self.session, self.base = session, base
+        self.console = console or Console()
+        self.render = Renderer(self.console)
+        self.temperature, self.max_new_tokens = temperature, max_new_tokens
+        self.raw, self.show_logits = raw, show_logits
+        self.compare_on = compare and base is not None
+        self.system, self.seed, self.stream = system, seed, stream
+        self.messages: list[dict] = []
+        self.last_prompt: str | None = None   # exact text last fed to the model
+        self.registry = REGISTRY
+        self.reset()
+
+    # ------------------------------------------------------------------ state
+    def reset(self) -> None:
+        self.messages = [{"role": "system", "content": self.system}] if self.system else []
+
+    def require_prompt(self) -> str:
+        if self.last_prompt is None:
+            raise CommandError("no prompt yet — say something first")
+        return self.last_prompt
+
+    def sampling(self) -> dict:
+        return dict(max_new_tokens=self.max_new_tokens, temperature=self.temperature, seed=self.seed)
+
+    # ------------------------------------------------------------------ turns
+    def _gen(self, sess: ModelSession, msgs: list[dict], stream: bool):
+        on_text = None
+        if stream:
+            on_text = lambda s: self.console.print(s, end="", markup=False, highlight=False)  # noqa: E731
+        if self.raw:
+            return sess.generate(sess.format_chat(msgs, raw=True), on_text=on_text, **self.sampling())
+        return sess.chat(msgs, on_text=on_text, **self.sampling())
+
+    def ask(self, text: str) -> str:
+        """One user turn: generate (and stream) the reply, then auto-analysis. Returns the reply."""
+        msgs = [{"role": "user", "content": text}] if self.raw else self.messages + [{"role": "user", "content": text}]
+        self.last_prompt = self.session.format_chat(msgs, raw=self.raw)
+        streaming = self.stream and not self.compare_on
+        if self.compare_on:
+            self.console.rule(f"[cyan]{self.session.name}")
+        g = self._gen(self.session, msgs, streaming)
+        if streaming:
+            self.console.print()
+        else:
+            self.console.print(g.text, markup=False, highlight=False)
+        if self.compare_on:
+            self.console.rule(f"[magenta]{self.base.name}")
+            gb = self._gen(self.base, msgs, False)
+            self.console.print(gb.text, markup=False, highlight=False)
+            self.console.rule()
+            c = self.session.compare(self.base, self.last_prompt, k=self.show_logits or 5)
+            self.render.comparison(c, self.session.name, self.base.name)
+        if self.show_logits:
+            self.render.topk(self.session.next_token_topk(self.last_prompt, self.show_logits),
+                             title=f"next token · {self.session.name}")
+        if not self.raw:
+            self.messages = msgs + [{"role": "assistant", "content": g.text}]
+        return g.text
+
+    def handle(self, line: str):
+        """Process one input line (command or chat turn). Returns 'quit' to stop."""
+        line = line.strip()
+        if not line:
+            return None
+        if line.startswith("/"):
+            try:
+                return self.registry.dispatch(self, line)
+            except CommandError as e:
+                self.console.print(f"[red]{e}[/]")
+                return None
+        self.ask(line)
+        return None
+
+    # ------------------------------------------------------------------ REPL
+    def _reader(self):
+        if sys.stdin.isatty():
+            try:
+                from prompt_toolkit import PromptSession
+                from prompt_toolkit.completion import WordCompleter
+                from prompt_toolkit.history import FileHistory
+                hist = cache_root() / "chat_history"
+                hist.parent.mkdir(parents=True, exist_ok=True)
+                ps = PromptSession(history=FileHistory(str(hist)),
+                                   completer=WordCompleter(self.registry.names(), sentence=True))
+                return lambda: ps.prompt("you> ")
+            except (ImportError, OSError):  # no prompt_toolkit, or an unusable terminal / history file
+                pass
+        return lambda: input("you> ")
+
+    def repl(self) -> int:
+        s = self.session
+        mode = "raw completion" if self.raw else ("chat template" if s.has_chat_template else "User/Assistant transcript")
+        info = (f"[bold]{s.name}[/] ({s.lm.source_kind}"
+                + (f", base {s.lm.base_name}" if s.lm.base_name else "") + f") · {mode} · temp {self.temperature}")
+        if self.base is not None:
+            info += f"\ncompare: {self.base.name} ({'on' if self.compare_on else 'off'}, /compare toggles)"
+        self.console.print(Panel(info + "\n/help for commands", title="llmf chat", expand=False))
+        read = self._reader()
+        while True:
+            try:
+                line = read()
+            except (EOFError, KeyboardInterrupt):
+                self.console.print()
+                return 0
+            try:
+                if self.handle(line) == "quit":
+                    return 0
+            except KeyboardInterrupt:
+                self.console.print("[yellow]interrupted[/]")
+
+
+def run(a: argparse.Namespace, console: Console | None = None) -> int:
+    console = console or Console()
+    with console.status(f"loading {a.weights} ..."):
+        session = ModelSession(resolve_weights(a.weights, base=a.base, device=a.device, dtype=a.dtype))
+        base = None
+        if a.base:
+            base = ModelSession(resolve_weights(a.base, device=a.device, dtype=a.dtype))
+    app = ChatApp(session, base, console=console, temperature=a.temperature,
+                  max_new_tokens=a.max_new_tokens, raw=a.raw, show_logits=a.show_logits,
+                  compare=base is not None, system=a.system, seed=a.seed,
+                  stream=not a.no_stream and a.prompt is None)
+    if a.prompt is not None:
+        app.ask(a.prompt)
+        return 0
+    return app.repl()
+
+
+def main(argv: list[str] | None = None) -> int:
+    return run(build_parser().parse_args(argv))
+
+
+if __name__ == "__main__":  # pragma: no cover
+    raise SystemExit(main())

llm_forensics/chat/args.py

@@ -0,0 +1,32 @@
+"""Argument definitions for `llmf chat` (no torch import, so `llmf --help` stays fast)."""
+from __future__ import annotations
+
+import argparse
+
+from llm_forensics.cli import DTYPE_CHOICES
+
+
+def build_parser(prog: str = "llmf chat") -> argparse.ArgumentParser:
+    ap = argparse.ArgumentParser(
+        prog=prog, description="Chat with a model and inspect its next-token predictions.\n\n"
+        "WEIGHTS can be an HF id, local dir, run dir, .safetensors file (with --base) or gs:// path.",
+        formatter_class=argparse.RawDescriptionHelpFormatter)
+    add_arguments(ap)
+    return ap
+
+
+def add_arguments(ap: argparse.ArgumentParser) -> None:
+    ap.add_argument("weights", metavar="WEIGHTS", help="the model to chat with")
+    ap.add_argument("--base", help="base (pre-edit) weights: shown side by side with KL, toggled with /compare; "
+                                   "also supplies config and tokenizer for a bare .safetensors file")
+    ap.add_argument("--prompt", help="answer this prompt and exit")
+    ap.add_argument("--max-new-tokens", type=int, default=128, help="tokens to generate per answer (default 128)")
+    ap.add_argument("--temperature", type=float, default=0.0, help="sampling temperature; 0 = greedy (default)")
+    ap.add_argument("--seed", type=int, default=None, help="random seed for sampling")
+    ap.add_argument("--show-logits", type=int, default=0, metavar="K",
+                    help="after each answer, show the top-K next tokens for the prompt")
+    ap.add_argument("--raw", action="store_true", help="plain text completion, without a chat template")
+    ap.add_argument("--system", help="system message")
+    ap.add_argument("--device", default=None, help="cpu, cuda or mps (default: auto)")
+    ap.add_argument("--dtype", default=None, choices=DTYPE_CHOICES, help="default: float32 on CPU, bfloat16 on GPU")
+    ap.add_argument("--no-stream", action="store_true", help="print each answer when it is complete")

llm_forensics/chat/commands.py

@@ -0,0 +1,176 @@
+"""Slash-command registry for `llmf chat`.
+
+Adding an analysis command is one function + decorator:
+
+    @command("entropy", help="entropy of the next-token distribution", usage="/entropy")
+    def _entropy(app, args: str) -> None:
+        t = app.session.next_token_topk(app.require_prompt(), k=1)
+        app.console.print(f"H = {t.entropy:.3f} nats")
+
+The function receives the ChatApp and the raw argument string. Raise CommandError for a
+user-facing error message. Return the string "quit" to leave the REPL.
+"""
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Callable
+
+if TYPE_CHECKING:  # pragma: no cover
+    from .app import ChatApp
+
+Handler = Callable[["ChatApp", str], object]
+
+
+class CommandError(Exception):
+    """Message shown to the user; the REPL keeps running."""
+
+
+@dataclass
+class Command:
+    name: str
+    func: Handler
+    help: str = ""
+    usage: str = ""
+    aliases: tuple[str, ...] = field(default_factory=tuple)
+
+
+class CommandRegistry:
+    def __init__(self):
+        self._cmds: dict[str, Command] = {}
+        self._alias: dict[str, str] = {}
+
+    def register(self, name: str, help: str = "", usage: str | None = None,
+                 aliases: tuple[str, ...] | list[str] = ()) -> Callable[[Handler], Handler]:
+        def deco(func: Handler) -> Handler:
+            if name in self._cmds or name in self._alias:
+                raise ValueError(f"command /{name} already registered")
+            self._cmds[name] = Command(name, func, help or (func.__doc__ or "").strip().split("\n")[0],
+                                       usage or f"/{name}", tuple(aliases))
+            for a in aliases:
+                self._alias[a] = name
+            return func
+        return deco
+
+    def get(self, name: str) -> Command | None:
+        name = name.lstrip("/")
+        return self._cmds.get(self._alias.get(name, name))
+
+    def __contains__(self, name: str) -> bool:
+        return self.get(name) is not None
+
+    def __iter__(self):
+        return iter(self._cmds.values())
+
+    def names(self) -> list[str]:
+        return ["/" + n for n in self._cmds] + ["/" + a for a in self._alias]
+
+    def dispatch(self, app: "ChatApp", line: str):
+        """Run a '/name args' line. Raises CommandError for unknown commands."""
+        line = line.strip()
+        if not line.startswith("/"):
+            raise CommandError("not a command")
+        name, _, args = line[1:].partition(" ")
+        cmd = self.get(name)
+        if cmd is None:
+            raise CommandError(f"unknown command /{name} (try /help)")
+        return cmd.func(app, args.strip())
+
+
+REGISTRY = CommandRegistry()
+command = REGISTRY.register
+
+
+def parse_ints(args: str) -> list[int]:
+    """'3 5 7' / '3,5,7' / '0-11' / '0-11:2' -> list of ints."""
+    out: list[int] = []
+    for tok in args.replace(",", " ").split():
+        rng, _, step = tok.partition(":")
+        m = re.fullmatch(r"(-?\d+)-(-?\d+)", rng)
+        if m:
+            out += list(range(int(m[1]), int(m[2]) + 1, int(step or 1)))
+        else:
+            out.append(int(rng))
+    return out
+
+
+# --------------------------------------------------------------------------- built-ins
+
+@command("help", help="list commands", aliases=("h", "?"))
+def _help(app, args):
+    app.render.help(REGISTRY)
+
+
+@command("quit", help="exit", aliases=("exit", "q"))
+def _quit(app, args):
+    return "quit"
+
+
+@command("logits", help="top-k next-token table for the last prompt", usage="/logits [k]")
+def _logits(app, args):
+    k = int(args) if args else app.show_logits or 10
+    prompt = app.require_prompt()
+    app.render.topk(app.session.next_token_topk(prompt, k), title=f"next token · {app.session.name}")
+    if app.compare_on and app.base is not None:
+        app.render.topk(app.base.next_token_topk(prompt, k), title=f"next token · {app.base.name}")
+
+
+@command("lens", help="logit lens: top tokens per layer at the last prompt position",
+         usage="/lens [layers e.g. 0-11:2 or 3 6 9] [k=5]")
+def _lens(app, args):
+    k = 5
+    parts = []
+    for tok in args.split():
+        if tok.startswith("k="):
+            k = int(tok[2:])
+        else:
+            parts.append(tok)
+    try:
+        layers = parse_ints(" ".join(parts)) if parts else None
+    except ValueError as e:
+        raise CommandError(f"bad layer spec {args!r}: {e}") from None
+    app.render.lens(app.session.logit_lens(app.require_prompt(), layers, k=k), title=app.session.name)
+
+
+@command("compare", help="toggle side-by-side completions against the --base model, with KL")
+def _compare(app, args):
+    if app.base is None:
+        raise CommandError("no base model; start with --base <weights>")
+    app.compare_on = not app.compare_on if args not in ("on", "off") else args == "on"
+    app.console.print(f"compare: [bold]{'on' if app.compare_on else 'off'}[/] ({app.base.name})")
+
+
+@command("temp", help="set sampling temperature (0 = greedy)", usage="/temp X")
+def _temp(app, args):
+    try:
+        app.temperature = float(args)
+    except ValueError:
+        raise CommandError("usage: /temp X") from None
+    app.console.print(f"temperature = {app.temperature}")
+
+
+@command("max", help="set max new tokens", usage="/max N")
+def _max(app, args):
+    try:
+        app.max_new_tokens = int(args)
+    except ValueError:
+        raise CommandError("usage: /max N") from None
+    app.console.print(f"max_new_tokens = {app.max_new_tokens}")
+
+
+@command("reset", help="clear the conversation history")
+def _reset(app, args):
+    app.reset()
+    app.console.print("history cleared")
+
+
+@command("raw", help="toggle raw completion mode (no chat template / transcript, no history)")
+def _raw(app, args):
+    app.raw = not app.raw if args not in ("on", "off") else args == "on"
+    app.reset()
+    app.console.print(f"raw mode: [bold]{'on' if app.raw else 'off'}[/] (history cleared)")
+
+
+@command("prompt", help="show the exact text of the last prompt fed to the model")
+def _prompt(app, args):
+    app.console.print(repr(app.require_prompt()), markup=False, highlight=False)

llm_forensics/chat/render.py

@@ -0,0 +1,74 @@
+"""Rich rendering of core.session results (tables with probability bars)."""
+from __future__ import annotations
+
+from rich.console import Console
+from rich.markup import escape
+from rich.table import Table
+
+from llm_forensics.core.session import Comparison, LogitLens, TopK
+
+
+def bar(p: float, width: int = 24) -> str:
+    """Unicode probability bar, e.g. bar(0.5, 8) -> '████    '."""
+    p = min(max(p, 0.0), 1.0)
+    eighths = int(round(p * width * 8))
+    full, rem = divmod(eighths, 8)
+    s = "█" * full + (" ▏▎▍▌▋▊▉"[rem] if rem else "")
+    return s.ljust(width)
+
+
+def tok_repr(t: str) -> str:
+    return repr(t)[1:-1] if t.strip() == t and t else repr(t)
+
+
+class Renderer:
+    def __init__(self, console: Console):
+        self.console = console
+
+    def topk(self, t: TopK, title: str = "next token") -> Table:
+        tab = Table(title=f"{escape(title)}  (H={t.entropy:.2f} nats)", title_justify="left", expand=False)
+        tab.add_column("#", justify="right", style="dim")
+        tab.add_column("token", style="cyan", no_wrap=True)
+        tab.add_column("id", justify="right", style="dim")
+        tab.add_column("prob", justify="right")
+        tab.add_column("", style="green", no_wrap=True)
+        tab.add_column("logit", justify="right", style="dim")
+        for i, (tok, tid, p, lg) in enumerate(zip(t.tokens, t.ids, t.probs, t.logits), 1):
+            tab.add_row(str(i), escape(tok_repr(tok)), str(tid), f"{p:.4f}", bar(p), f"{lg:.2f}")
+        self.console.print(tab)
+        return tab
+
+    def lens(self, ll: LogitLens, title: str = "") -> Table:
+        k = len(ll.layers[0].tokens) if ll.layers else 0
+        tab = Table(title=f"logit lens {escape(title)} @pos {ll.position}", title_justify="left")
+        tab.add_column("layer", justify="right", style="dim")
+        for j in range(k):
+            tab.add_column(f"top{j + 1}", no_wrap=True)
+        for L in ll.layers:
+            tab.add_row(str(L.layer), *[f"[cyan]{escape(tok_repr(t))}[/] [dim]{p:.2f}[/]"
+                                        for t, p in zip(L.tokens, L.probs)])
+        self.console.print(tab)
+        return tab
+
+    def comparison(self, c: Comparison, a_name: str, b_name: str) -> Table:
+        tab = Table(title=f"KL({escape(a_name)}‖{escape(b_name)})={c.kl:.4f}  KL(rev)={c.kl_reverse:.4f}  JS={c.js:.4f}"
+                          f"  top1 {'agree' if c.top1_agree else '[red]DIFFER[/]'}", title_justify="left")
+        tab.add_column("#", justify="right", style="dim")
+        tab.add_column(escape(a_name), style="cyan", no_wrap=True)
+        tab.add_column("p", justify="right")
+        tab.add_column(escape(b_name), style="magenta", no_wrap=True)
+        tab.add_column("p", justify="right")
+        for i in range(len(c.self_topk.ids)):
+            tab.add_row(str(i + 1), escape(tok_repr(c.self_topk.tokens[i])), f"{c.self_topk.probs[i]:.4f}",
+                        escape(tok_repr(c.other_topk.tokens[i])), f"{c.other_topk.probs[i]:.4f}")
+        self.console.print(tab)
+        return tab
+
+    def help(self, registry) -> None:
+        tab = Table(title="commands", title_justify="left", show_header=False, box=None)
+        tab.add_column(style="bold cyan", no_wrap=True)
+        tab.add_column()
+        for c in registry:
+            al = f" [dim](/{', /'.join(c.aliases)})[/]" if c.aliases else ""
+            tab.add_row(escape(c.usage), escape(c.help) + al)
+        self.console.print(tab)

llm_forensics/cli.py

@@ -0,0 +1,57 @@
+"""`llmf`: detect and explain edits hidden in LLM weights.
+
+<weights> is a Hugging Face id, a local HF directory, an `llmf edit` run directory, a
+.safetensors file (with --base) or a gs:// path. `llmf <command> -h` shows a command's options.
+"""
+from __future__ import annotations
+
+import argparse
+import importlib
+import sys
+
+from llm_forensics import BUILD
+
+DTYPE_CHOICES = ("float32", "bfloat16", "float16")   # --dtype in every command
+
+# command -> (module with main(argv), usage line, help)
+COMMANDS = {
+    "models": ("llm_forensics.models", "llmf models {register,list}",
+               "register existing base-model directories for reuse"),
+    "identify": ("llm_forensics.identify.cli", "llmf identify <weights> [--top-k K]",
+                 "identify the base/ancestor checkpoint of unknown weights"),
+    "scan": ("llm_forensics.scan.cli", "llmf scan <weights> [--base <base>]",
+             "scan weights for hidden edits"),
+    "edit": ("llm_forensics.edit.run", "llmf edit <config.yaml>",
+             "apply a ROME or MEMIT edit described by a YAML config"),
+    "chat": ("llm_forensics.chat.app", "llmf chat <weights>",
+             "chat with a model and inspect its next-token predictions"),
+}
+
+
+def build_parser() -> argparse.ArgumentParser:
+    usage = "\n".join(f"    {u:<38} {h}" for _, u, h in COMMANDS.values())
+    ap = argparse.ArgumentParser(prog="llmf", description=f"{__doc__.splitlines()[0]}\n\n{usage}\n\n"
+                                 + "\n".join(__doc__.splitlines()[2:]),
+                                 formatter_class=argparse.RawDescriptionHelpFormatter)
+    ap.add_argument("--version", action="version", version=f"llmf {BUILD}")
+    sub = ap.add_subparsers(dest="cmd", metavar="{scan,edit,chat,identify,models}")
+    for name, (_, _, help_) in COMMANDS.items():
+        sub.add_parser(name, help=help_, add_help=False)
+    return ap
+
+
+def main(argv: list[str] | None = None) -> int:
+    argv = sys.argv[1:] if argv is None else list(argv)
+    if argv and argv[0] in COMMANDS:  # each command parses its own options; heavy imports happen here
+        mod = importlib.import_module(COMMANDS[argv[0]][0])
+        return int(mod.main(argv[1:]) or 0)
+    ap = build_parser()
+    a = ap.parse_args(argv)
+    if a.cmd is None:
+        ap.print_help()
+        return 2
+    return 0  # pragma: no cover
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

llm_forensics/core/__init__.py

@@ -0,0 +1,3 @@
+"""Shared core: weights loading (`loading`) and reusable analysis primitives (`session`)."""
+from .loading import LoadedModel, TensorReader, open_tensors, resolve_weights  # noqa: F401
+from .session import Comparison, Generation, LogitLens, MLPKeys, ModelSession, TopK  # noqa: F401