logogram 0.1.0__py3-none-any.whl

logogram/datasets.py

@@ -0,0 +1,159 @@
+"""Prompt datasets: JSONL files of clean/corrupt pairs with an answer and a distractor.
+
+Each line is one JSON object::
+
+    {"clean": "When Mary and John went to the store, John gave a drink to",
+     "corrupt": "When Mary and John went to the store, Mary gave a drink to",
+     "answer": " Mary", "distractor": " John",
+     "positions": {"IO": [5, 9], "S1": [14, 18], "S2": [38, 42]}}
+
+``positions`` is optional. It names character spans in the clean prompt; a label points to the
+last token that overlaps its span. ``id`` and ``meta`` are optional and carried through.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import re
+from pathlib import Path
+from typing import Any
+
+from pydantic import BaseModel, ConfigDict, Field, ValidationError, field_validator
+
+from logogram.fileio import write_text_atomic
+
+DATASET_NAME_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]{0,99}$")
+
+
+class DatasetError(ValueError):
+    """A dataset file can't be read. The message says where and how to fix it."""
+
+
+class PromptRecord(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    clean: str = Field(min_length=1)
+    corrupt: str = Field(min_length=1)
+    answer: str = Field(min_length=1)
+    distractor: str = Field(min_length=1)
+    positions: dict[str, tuple[int, int]] | None = None
+    id: str | None = None
+    meta: dict[str, Any] | None = None
+
+    @field_validator("positions")
+    @classmethod
+    def _spans_ordered(
+        cls, value: dict[str, tuple[int, int]] | None
+    ) -> dict[str, tuple[int, int]] | None:
+        if value is None:
+            return value
+        for label, (start, end) in value.items():
+            if not label or not label.strip():
+                raise ValueError("position labels must be non-empty")
+            if label in ("all", "last") or re.fullmatch(r"-?\d+", label):
+                raise ValueError(
+                    f"{label!r} can't be a position label: 'all', 'last' and numbers already "
+                    "name positions"
+                )
+            if start < 0 or end <= start:
+                raise ValueError(f"position {label!r} must be a span [start, end) with start < end")
+        return value
+
+
+def parse_jsonl(text: str, *, source: str = "dataset") -> list[PromptRecord]:
+    records: list[PromptRecord] = []
+    for line_no, line in enumerate(text.splitlines(), start=1):
+        if not line.strip():
+            continue
+        try:
+            raw = json.loads(line)
+        except json.JSONDecodeError as exc:
+            raise DatasetError(
+                f"{source}, line {line_no}: not valid JSON ({exc.msg}). "
+                "Each line must be one JSON object."
+            ) from exc
+        if not isinstance(raw, dict):
+            raise DatasetError(f"{source}, line {line_no}: expected a JSON object.")
+        missing = [k for k in ("clean", "corrupt", "answer", "distractor") if k not in raw]
+        if missing:
+            raise DatasetError(
+                f"{source}, line {line_no}: missing {', '.join(repr(m) for m in missing)}. "
+                "Each line needs clean, corrupt, answer and distractor."
+            )
+        if isinstance(raw.get("positions"), dict):
+            raw["positions"] = {
+                k: tuple(v) if isinstance(v, list) else v for k, v in raw["positions"].items()
+            }
+        try:
+            record = PromptRecord.model_validate(raw)
+        except ValidationError as exc:
+            first = exc.errors()[0]
+            where = ".".join(str(p) for p in first["loc"])
+            raise DatasetError(f"{source}, line {line_no}: {where}: {first['msg']}.") from exc
+        _check_spans(record, line_no, source)
+        records.append(record)
+    if not records:
+        raise DatasetError(f"{source} has no prompts. Add one JSON object per line.")
+    return records
+
+
+def _check_spans(record: PromptRecord, line_no: int, source: str) -> None:
+    if not record.positions:
+        return
+    for label, (_start, end) in record.positions.items():
+        if end > len(record.clean):
+            raise DatasetError(
+                f"{source}, line {line_no}: position {label!r} ends at character {end}, but "
+                f"the clean prompt has {len(record.clean)} characters."
+            )
+
+
+def load_dataset(path: Path) -> list[PromptRecord]:
+    if not path.exists():
+        raise DatasetError(f"Dataset not found: {path.name}. Check the path in the spec.")
+    if not path.is_file():  # a folder, device or pipe could fail, block or never end
+        raise DatasetError(f"{path.name} isn't a regular file.")
+    try:
+        text = path.read_text(encoding="utf-8")
+    except UnicodeDecodeError as exc:
+        raise DatasetError(
+            f"{path.name} isn't UTF-8 text (byte {exc.start} can't be decoded). Save it as UTF-8."
+        ) from exc
+    except OSError as exc:
+        raise DatasetError(f"{path.name} can't be read: {exc.strerror or exc}.") from exc
+    return parse_jsonl(text, source=path.name)
+
+
+def file_sha256(path: Path) -> str:
+    digest = hashlib.sha256()
+    with path.open("rb") as fh:
+        for chunk in iter(lambda: fh.read(1 << 20), b""):
+            digest.update(chunk)
+    return digest.hexdigest()
+
+
+def to_jsonl(records: list[PromptRecord]) -> str:
+    lines = []
+    for record in records:
+        data = record.model_dump(mode="json", exclude_none=True)
+        if "positions" in data:
+            data["positions"] = {k: list(v) for k, v in data["positions"].items()}
+        lines.append(json.dumps(data, ensure_ascii=False))
+    return "\n".join(lines) + "\n"
+
+
+def write_dataset(path: Path, records: list[PromptRecord]) -> None:
+    """Write a dataset file. Callers in a project get ``path`` from ``Project.dataset_file``."""
+    path.parent.mkdir(parents=True, exist_ok=True)
+    write_text_atomic(path, to_jsonl(records))
+
+
+def check_dataset_name(name: str) -> str:
+    stem = name[:-6] if name.endswith(".jsonl") else name
+    if not DATASET_NAME_RE.match(stem):
+        raise DatasetError(
+            "Dataset names can use letters, digits, '.', '_' and '-', and must start with a "
+            "letter or digit."
+        )
+    return stem + ".jsonl"

logogram/direct.py

@@ -0,0 +1,193 @@
+"""Direct logit attribution: what each component writes straight into the logit difference.
+
+At the last position the residual stream is the embeddings plus every attention and MLP output.
+With the final normalization's scale held at its value in the run, the logit difference is an
+affine function of that sum, so it splits into one term per component plus a constant from biases.
+A component's term is its *direct* effect: it leaves out everything the component does through
+later components, which patching measures. The backend computes the terms; this module checks that
+the split is valid for the loaded model and the chosen sites, and turns the terms into a result.
+"""
+
+from __future__ import annotations
+
+import threading
+from collections.abc import Callable
+from typing import Any
+
+import numpy as np
+
+from logogram.backends.base import Cancelled, ModelBackend, ModelInfo
+from logogram.engine import (
+    EngineError,
+    EngineResult,
+    LayerFn,
+    ProgressFn,
+    _answer_tensors,
+    _chunks,
+    compute_baselines,
+)
+from logogram.prompts import PreparedPrompt, group_by_length
+from logogram.sites import ResolvedSite, ScopeError, expand_scope, resolve_position
+from logogram.spec import DirectLogitAttribution, Spec
+
+DIRECT_KINDS = ("head", "attn_out", "mlp_out")
+
+
+def check_direct_sites(
+    sites: list[ResolvedSite], info: ModelInfo, prompts: list[PreparedPrompt]
+) -> None:
+    """Refuse sites, or models, for which a direct effect isn't defined."""
+    extra = info.extra
+    if not str(extra.get("block_structure", "")).startswith(("sequential", "parallel")):
+        raise ScopeError(
+            "When this model loaded, its residual stream couldn't be verified to be the sum of its "
+            "attention and MLP outputs, so the logit difference can't be split among them."
+        )
+    if extra.get("logit_soft_cap"):
+        raise ScopeError(
+            "This model soft-caps its logits (tanh), so the logit difference isn't a sum of "
+            "direct effects. Use activation patching for this model."
+        )
+    for rs in sites:
+        if rs.kind not in DIRECT_KINDS:
+            raise ScopeError(
+                f"{rs.label} is a state of the residual stream, not something written into it. "
+                "Direct logit attribution splits the logit difference among heads, attention "
+                "outputs and MLP outputs; choose those."
+            )
+    if any(rs.kind == "head" for rs in sites):
+        checks = extra.get("checks") or {}
+        heads = checks.get("heads")
+        if heads is None or heads > checks.get("tolerance", 0):
+            raise ScopeError(
+                "In this model the attention output isn't the sum of its heads' outputs (it is "
+                "normalized after they are combined), so a single head has no direct effect of its "
+                "own. Use attention and MLP outputs per layer."
+            )
+    checked: set[str] = set()
+    for rs in sites:
+        key = rs.site.position.model_dump_json()
+        if key in checked:
+            continue
+        checked.add(key)
+        for prompt in prompts:
+            if resolve_position(rs.site.position, prompt) != prompt.length - 1:
+                raise ScopeError(
+                    "Direct effects are read at the last token, where the logit difference is "
+                    f"measured, but {rs.label} is at another position in prompt {prompt.index}. "
+                    "Set the position to the last token."
+                )
+
+
+def run_direct_effects(
+    spec: Spec,
+    backend: ModelBackend,
+    prompts: list[PreparedPrompt],
+    *,
+    on_progress: ProgressFn | None = None,
+    on_layer: LayerFn | None = None,
+    cancel: threading.Event | None = None,
+    on_start: Callable[[list[ResolvedSite], dict[str, Any]], None] | None = None,
+) -> EngineResult:
+    exp = spec.experiment
+    assert isinstance(exp, DirectLogitAttribution)
+    info = backend.info
+    batch_size = spec.execution.batch_size
+    sites, layout = expand_scope(spec, info, prompts)
+    check_direct_sites(sites, info, prompts)
+    if on_start is not None:
+        on_start(sites, layout)
+    groups = group_by_length(prompts)
+    baselines = compute_baselines(backend, prompts, groups, batch_size, cancel)
+    which = exp.prompts
+    n, n_layers = len(prompts), info.n_layers
+    heads = any(rs.kind == "head" for rs in sites)
+
+    terms: dict[str, np.ndarray] = {
+        "embed": np.zeros(n),
+        "attn_out": np.zeros((n, n_layers)),
+        "mlp_out": np.zeros((n, n_layers)),
+        "logit_diff": np.zeros(n),
+        "remainder": np.zeros(n),
+    }
+    if heads:
+        terms["head"] = np.zeros((n, n_layers, info.n_heads))
+    done = 0
+    for group in groups:
+        tokens = group.clean if which == "clean" else group.corrupt
+        for sl in _chunks(len(group.members), batch_size):
+            if cancel is not None and cancel.is_set():
+                raise Cancelled()
+            idx = group.members[sl]
+            out = backend.direct_effects(tokens[sl], *_answer_tensors(prompts, idx), heads)
+            for key, values in out.items():
+                terms[key][idx] = values.numpy()
+            done += len(idx)
+            if on_progress is not None:
+                on_progress(done, n, n_layers - 1)
+
+    gap = baselines.ld(which)
+    warnings: list[str] = []
+    mean_gap = float(gap.mean())
+    if spec.metric.normalization == "dataset_gap" and abs(mean_gap) < 1e-3:
+        raise EngineError(
+            f"The {which} prompts' mean logit difference is almost zero ({mean_gap:.4f}), so a "
+            "share of it is undefined. Normalize by each prompt's own logit difference, or check "
+            "the baseline."
+        )
+    if spec.metric.normalization == "prompt_gap":
+        zero = [p.index for p, g in zip(prompts, gap, strict=True) if abs(g) < 1e-6]
+        if zero:
+            raise EngineError(
+                f"Prompt(s) {', '.join(map(str, zero[:5]))} have a logit difference of zero, so a "
+                "share of it is undefined. Normalize by the dataset's mean, or fix those prompts."
+            )
+        small = int((np.abs(gap) < 0.1).sum())
+        if small:
+            warnings.append(
+                f"{small} prompt(s) have a logit difference below 0.1, so their per-prompt shares "
+                "are unstable."
+            )
+    # The split is of the model's own logit difference; the two ways of reading it out must agree.
+    drift = float(np.abs(terms["logit_diff"] - gap).max()) if n else 0.0
+    if drift > 1e-3 * max(1.0, float(np.abs(gap).max())):
+        warnings.append(
+            f"The decomposed logit difference differs from the measured one by up to {drift:.2g}, "
+            "more than rounding explains."
+        )
+
+    delta = np.zeros((len(sites), n))
+    for rs in sites:
+        if rs.kind == "head":
+            delta[rs.index] = terms["head"][:, rs.layer, rs.head]
+        else:
+            delta[rs.index] = terms[rs.kind][:, rs.layer]
+    nan = np.full((len(sites), n), np.nan)
+    result = EngineResult(
+        sites=sites,
+        layout=layout,
+        prompts=prompts,
+        baselines=baselines,
+        receiver=which,
+        reference=which,
+        patched_ld=nan,
+        patched_prob=nan.copy(),
+        warnings=warnings,
+        measure="attribution",
+        delta=delta,
+        gap=gap,
+        extra={
+            "direct": {
+                "prompts": which,
+                "logit_diff": float(terms["logit_diff"].mean()),
+                "embeddings": float(terms["embed"].mean()),
+                "attention": float(terms["attn_out"].sum(1).mean()),
+                "mlp": float(terms["mlp_out"].sum(1).mean()),
+                "biases": float(terms["remainder"].mean()),
+            }
+        },
+    )
+    if on_layer is not None:
+        for layer in sorted({rs.layer for rs in sites}):
+            on_layer(layer, [rs.index for rs in sites if rs.layer == layer], result)
+    return result