logogram 0.1.0__py3-none-any.whl

logogram/prompts.py

@@ -0,0 +1,204 @@
+"""Tokenize prompt pairs for a model and check that they can be used in an experiment.
+
+A usable pair has clean and corrupt prompts of the same token length (so positions line up), and
+an answer and distractor that are each a single token.
+"""
+
+from __future__ import annotations
+
+from collections import defaultdict
+from dataclasses import dataclass, field
+
+import torch
+
+from logogram.backends.base import ModelBackend, Tokenized
+from logogram.datasets import PromptRecord
+
+
+@dataclass
+class PromptIssue:
+    index: int
+    kind: str
+    message: str
+
+    def to_dict(self) -> dict[str, object]:
+        return {"index": self.index, "kind": self.kind, "message": self.message}
+
+
+@dataclass
+class PreparedPrompt:
+    index: int
+    record: PromptRecord
+    clean: Tokenized
+    corrupt: Tokenized
+    answer_id: int
+    distractor_id: int
+    labels: dict[str, int] = field(default_factory=dict)
+
+    @property
+    def length(self) -> int:
+        return len(self.clean.ids)
+
+    def differing_positions(self) -> list[int]:
+        return [
+            i
+            for i, (a, b) in enumerate(zip(self.clean.ids, self.corrupt.ids, strict=True))
+            if a != b
+        ]
+
+
+class PromptError(ValueError):
+    def __init__(self, issues: list[PromptIssue]):
+        self.issues = issues
+        shown = "; ".join(f"prompt {i.index}: {i.message}" for i in issues[:3])
+        more = f" (and {len(issues) - 3} more)" if len(issues) > 3 else ""
+        super().__init__(f"{len(issues)} prompt(s) can't be used. {shown}{more}")
+
+
+def _show_tokens(backend: ModelBackend, text: str) -> str:
+    pieces = backend.tokenize(text, prepend_bos=False).tokens
+    return " + ".join(repr(p) for p in pieces)
+
+
+def label_position(tokenized: Tokenized, span: tuple[int, int]) -> int | None:
+    """The last token whose character span overlaps ``span``."""
+    start, end = span
+    found = None
+    for i, (a, b) in enumerate(tokenized.offsets):
+        if b > a and a < end and b > start:
+            found = i
+    return found
+
+
+def prepare_prompt(
+    backend: ModelBackend, record: PromptRecord, index: int, prepend_bos: bool
+) -> tuple[PreparedPrompt | None, list[PromptIssue], Tokenized, Tokenized]:
+    issues: list[PromptIssue] = []
+    clean = backend.tokenize(record.clean, prepend_bos)
+    corrupt = backend.tokenize(record.corrupt, prepend_bos)
+    if len(clean.ids) != len(corrupt.ids):
+        issues.append(
+            PromptIssue(
+                index,
+                "length_mismatch",
+                f"clean has {len(clean.ids)} tokens and corrupt has {len(corrupt.ids)}, so "
+                "positions can't be aligned. Make both prompts tokenize to the same length.",
+            )
+        )
+    answer_id = backend.single_token_id(record.answer)
+    if answer_id is None:
+        issues.append(
+            PromptIssue(
+                index,
+                "answer_tokens",
+                f"the answer {record.answer!r} is several tokens "
+                f"({_show_tokens(backend, record.answer)}). Use a single-token answer, often "
+                "with a leading space.",
+            )
+        )
+    distractor_id = backend.single_token_id(record.distractor)
+    if distractor_id is None:
+        issues.append(
+            PromptIssue(
+                index,
+                "distractor_tokens",
+                f"the distractor {record.distractor!r} is several tokens "
+                f"({_show_tokens(backend, record.distractor)}). Use a single-token distractor.",
+            )
+        )
+    if clean.ids == corrupt.ids:
+        issues.append(
+            PromptIssue(
+                index,
+                "identical",
+                "the clean and corrupt prompts are the same, so there is nothing to patch.",
+            )
+        )
+    if answer_id is not None and answer_id == distractor_id:
+        issues.append(
+            PromptIssue(index, "same_answer", "the answer and distractor are the same token.")
+        )
+    n_ctx = backend.info.n_ctx
+    if max(len(clean.ids), len(corrupt.ids)) > n_ctx:
+        issues.append(
+            PromptIssue(index, "too_long", f"the prompt is longer than the model's {n_ctx} tokens.")
+        )
+    labels: dict[str, int] = {}
+    for label, span in (record.positions or {}).items():
+        pos = label_position(clean, span)
+        if pos is None:
+            issues.append(
+                PromptIssue(
+                    index,
+                    "label",
+                    f"position {label!r} (characters {span[0]}–{span[1]}) doesn't cover any token.",
+                )
+            )
+        else:
+            labels[label] = pos
+    if issues:
+        return None, issues, clean, corrupt
+    assert answer_id is not None and distractor_id is not None
+    prepared = PreparedPrompt(
+        index=index,
+        record=record,
+        clean=clean,
+        corrupt=corrupt,
+        answer_id=answer_id,
+        distractor_id=distractor_id,
+        labels=labels,
+    )
+    return prepared, [], clean, corrupt
+
+
+def prepare_prompts(
+    backend: ModelBackend, records: list[PromptRecord], prepend_bos: bool
+) -> list[PreparedPrompt]:
+    prepared: list[PreparedPrompt] = []
+    issues: list[PromptIssue] = []
+    for i, record in enumerate(records):
+        p, prompt_issues, _, _ = prepare_prompt(backend, record, i, prepend_bos)
+        issues.extend(prompt_issues)
+        if p is not None:
+            prepared.append(p)
+    if issues:
+        raise PromptError(issues)
+    return prepared
+
+
+@dataclass
+class LengthGroup:
+    """Prompts of one token length. Batches never mix lengths, so no padding is needed."""
+
+    length: int
+    members: list[int]  # indices into the prepared prompt list
+    clean: torch.Tensor  # [n_g, length]
+    corrupt: torch.Tensor
+
+
+def group_by_length(prompts: list[PreparedPrompt]) -> list[LengthGroup]:
+    by_len: dict[int, list[int]] = defaultdict(list)
+    for i, p in enumerate(prompts):
+        by_len[p.length].append(i)
+    groups = []
+    for length in sorted(by_len):
+        members = by_len[length]
+        groups.append(
+            LengthGroup(
+                length=length,
+                members=members,
+                clean=torch.tensor([prompts[i].clean.ids for i in members], dtype=torch.long),
+                corrupt=torch.tensor([prompts[i].corrupt.ids for i in members], dtype=torch.long),
+            )
+        )
+    return groups
+
+
+def common_labels(prompts: list[PreparedPrompt]) -> list[str]:
+    """Labels present in every prompt, ordered by their mean token position."""
+    if not prompts:
+        return []
+    shared = set(prompts[0].labels)
+    for p in prompts[1:]:
+        shared &= set(p.labels)
+    return sorted(shared, key=lambda lab: (sum(p.labels[lab] for p in prompts) / len(prompts), lab))

logogram/research.py

@@ -0,0 +1,84 @@
+"""Portable research notes and named selections, with optimistic edit revisions."""
+
+from __future__ import annotations
+
+import uuid
+from typing import Literal
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from logogram.fileio import write_text_atomic
+from logogram.project import Project, ProjectError, now_iso
+from logogram.spec import ModelRef, Site
+
+
+class NoteInput(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    title: str = Field(min_length=1, max_length=200)
+    body: str = Field(default="", max_length=50_000)
+    model: ModelRef
+    sites: list[Site] = Field(min_length=1, max_length=512)
+    run_id: str | None = Field(default=None, pattern=r"^[A-Za-z0-9][A-Za-z0-9._-]*$")
+
+
+class ResearchNote(NoteInput):
+    id: str
+    revision: int = Field(ge=1)
+    created: str
+    updated: str
+
+
+class Notebook(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    logogram_research: Literal[1] = 1
+    notes: list[ResearchNote] = Field(default_factory=list)
+
+
+class NoteConflict(ValueError):
+    pass
+
+
+def read_notebook(project: Project) -> Notebook:
+    data = project.read_json(project.root / "research.json", optional=True)
+    return Notebook.model_validate(data) if data is not None else Notebook()
+
+
+def save_note(
+    project: Project, value: NoteInput, *, note_id: str | None = None, revision: int | None = None
+) -> ResearchNote:
+    book = read_notebook(project)
+    old = next((n for n in book.notes if n.id == note_id), None)
+    if note_id is not None and (old is None or old.revision != revision):
+        raise NoteConflict("This note changed in another tab. Reload the note before saving again.")
+    if value.run_id is not None:
+        project.readable(project.run_dir(value.run_id) / "spec.json")
+    now = now_iso()
+    note = ResearchNote(
+        **value.model_dump(),
+        id=old.id if old else uuid.uuid4().hex,
+        revision=old.revision + 1 if old else 1,
+        created=old.created if old else now,
+        updated=now,
+    )
+    book.notes = [note if n.id == note.id else n for n in book.notes]
+    if old is None:
+        book.notes.append(note)
+    write_text_atomic(
+        project.writable(project.root / "research.json"), book.model_dump_json(indent=2) + "\n"
+    )
+    return note
+
+
+def delete_note(project: Project, note_id: str, revision: int) -> None:
+    book = read_notebook(project)
+    note = next((n for n in book.notes if n.id == note_id), None)
+    if note is None:
+        raise ProjectError("This note no longer exists. Refresh the notebook.")
+    if note.revision != revision:
+        raise NoteConflict("This note changed in another tab. Reload it before removing it.")
+    book.notes = [n for n in book.notes if n.id != note_id]
+    write_text_atomic(
+        project.writable(project.root / "research.json"), book.model_dump_json(indent=2) + "\n"
+    )

logogram/results.py

@@ -0,0 +1,240 @@
+"""Turn engine output into the files a run leaves behind: summary.json and results.parquet."""
+
+from __future__ import annotations
+
+import math
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import pyarrow as pa
+import pyarrow.parquet as pq
+
+from logogram.engine import EngineResult
+from logogram.fileio import atomic_output
+from logogram.schema import Summary
+from logogram.spec import Spec, describe_experiment
+from logogram.stats import SiteStats, compute_site_stats, resample_counts
+
+SUMMARY_VERSION = 1
+
+
+def _f(x: Any) -> float | None:
+    """A JSON-safe float (None for NaN or infinity)."""
+    value = float(x)
+    return value if math.isfinite(value) else None
+
+
+def compute_stats(
+    spec: Spec,
+    result: EngineResult,
+    counts: np.ndarray | None = None,
+    indices: list[int] | None = None,
+) -> SiteStats:
+    """Statistics for all sites, or for ``indices`` only (rows then follow ``indices``)."""
+    n = result.patched_ld.shape[1]
+    if counts is None:
+        counts = resample_counts(n, spec.statistics.bootstrap, spec.statistics.seed)
+    rows = slice(None) if indices is None else np.asarray(indices, dtype=np.int64)
+    return compute_site_stats(
+        patched_ld=result.patched_ld[rows],
+        patched_prob=result.patched_prob[rows],
+        receiver_ld=result.receiver_ld,
+        source_ld=result.reference_ld,
+        receiver_prob=result.receiver_prob,
+        normalization=spec.metric.normalization,
+        counts=counts,
+        ci=spec.statistics.ci,
+        delta=None if result.delta is None else result.delta[rows],
+        gap=result.gap,
+    )
+
+
+def site_payload(
+    result: EngineResult, stats: SiteStats, indices: list[int]
+) -> list[dict[str, Any]]:
+    """Payload for sites ``indices``; row ``i`` of ``stats`` belongs to ``indices[i]``."""
+    out = []
+    for i, site_index in enumerate(indices):
+        site = result.sites[site_index]
+        out.append(
+            {
+                **site.to_dict(),
+                "n": stats.n,
+                "effect": {
+                    "mean": _f(stats.effect_mean[i]),
+                    "sd": _f(stats.effect_sd[i]),
+                    "lo": _f(stats.effect_lo[i]),
+                    "hi": _f(stats.effect_hi[i]),
+                },
+                "delta": {
+                    "mean": _f(stats.delta_mean[i]),
+                    "sd": _f(stats.delta_sd[i]),
+                    "lo": _f(stats.delta_lo[i]),
+                    "hi": _f(stats.delta_hi[i]),
+                },
+                "patched_logit_diff": _f(stats.patched_mean[i]),
+                "answer_prob": _f(stats.prob_mean[i]),
+                "answer_prob_delta": _f(stats.prob_delta_mean[i]),
+                "sign_flips": int(stats.sign_flips[i]),
+                "opposite_sign": int(stats.opposite_sign[i]),
+            }
+        )
+    return out
+
+
+def _group_stats(values: np.ndarray) -> dict[str, float | None]:
+    sd = float(values.std(ddof=1)) if len(values) > 1 else 0.0
+    return {"mean": _f(values.mean()), "sd": _f(sd)}
+
+
+def baseline_payload(result: EngineResult) -> dict[str, Any]:
+    b = result.baselines
+    gap = result.reference_ld - result.receiver_ld
+    return {
+        "clean": {
+            "logit_diff": _group_stats(b.clean_ld),
+            "answer_prob": _group_stats(b.clean_prob),
+            "prefers_answer": int((b.clean_ld > 0).sum()),
+        },
+        "corrupt": {
+            "logit_diff": _group_stats(b.corrupt_ld),
+            "answer_prob": _group_stats(b.corrupt_prob),
+            "prefers_answer": int((b.corrupt_ld > 0).sum()),
+        },
+        "gap": _group_stats(gap),
+    }
+
+
+def build_summary(
+    spec: Spec,
+    result: EngineResult,
+    stats: SiteStats,
+    run_id: str,
+    model: dict[str, Any] | None = None,
+) -> dict[str, Any]:
+    b = result.baselines
+    normalization = spec.metric.normalization
+    description = "logit(answer) − logit(distractor) at the last position"
+    if result.measure == "attribution":
+        description = (
+            "direct contribution to logit(answer) − logit(distractor) at the last position, with "
+            "the final normalization's scale held at its value in the run"
+        )
+        if normalization == "dataset_gap":
+            norm_text = (
+                f"contribution ÷ mean {result.receiver} logit difference; "
+                f"mean = {stats.denominator:.4f}"
+            )
+        else:
+            norm_text = f"contribution ÷ each prompt's own {result.receiver} logit difference"
+    elif normalization == "dataset_gap":
+        norm_text = (
+            f"(patched − {result.receiver}) ÷ mean({result.reference} − {result.receiver}) "
+            f"logit difference; mean gap = {stats.denominator:.4f}"
+        )
+    else:
+        norm_text = (
+            f"(patched − {result.receiver}) ÷ ({result.reference} − {result.receiver}) "
+            "logit difference, per prompt"
+        )
+    if "steering" in result.extra:
+        norm_text = norm_text.replace("patched", "steered")
+    if result.measure == "estimate":
+        description = (
+            "first-order estimate of the change patching would cause in logit(answer) − "
+            "logit(distractor) at the last position: (source − receiver activation) · its "
+            "gradient at the receiver run"
+        )
+        norm_text = f"estimated {norm_text}"
+    summary = {
+        "logogram_summary": SUMMARY_VERSION,
+        "run_id": run_id,
+        "name": spec.name,
+        "description": describe_experiment(spec),
+        "model": model,
+        "n_prompts": len(result.prompts),
+        "n_sites": len(result.sites),
+        "layout": result.layout,
+        "receiver": result.receiver,
+        "reference": result.reference,
+        "measure": result.measure,
+        "metric": {
+            "kind": spec.metric.kind,
+            "normalization": normalization,
+            "denominator": _f(stats.denominator) if normalization == "dataset_gap" else None,
+            "description": description,
+            "normalized_effect": norm_text,
+        },
+        "statistics": {
+            "bootstrap": spec.statistics.bootstrap,
+            "ci": spec.statistics.ci,
+            "seed": spec.statistics.seed,
+            "method": "percentile bootstrap over prompts",
+        },
+        "baseline": baseline_payload(result),
+        "sites": site_payload(result, stats, list(range(len(result.sites)))),
+        "per_prompt": {
+            "clean_logit_diff": [_f(x) for x in b.clean_ld],
+            "corrupt_logit_diff": [_f(x) for x in b.corrupt_ld],
+            "clean_answer_prob": [_f(x) for x in b.clean_prob],
+            "corrupt_answer_prob": [_f(x) for x in b.corrupt_prob],
+        },
+        "donors": result.donors,
+        "warnings": result.warnings,
+        **({"direct": result.extra["direct"]} if "direct" in result.extra else {}),
+        **({"steering": result.extra["steering"]} if "steering" in result.extra else {}),
+        **({"features": result.extra["features"]} if "features" in result.extra else {}),
+    }
+    return Summary.model_validate(summary).model_dump(mode="json")
+
+
+def results_table(result: EngineResult, stats: SiteStats) -> pa.Table:
+    """Long format: one row per (site, prompt), with everything needed to recompute effects."""
+    n_sites, n = result.patched_ld.shape
+    site_idx = np.repeat(np.arange(n_sites, dtype=np.int32), n)
+    # The prompt's index in the dataset (a method may measure only some prompts, as steering
+    # measures the held-out ones).
+    prompt_idx = np.tile(np.array([p.index for p in result.prompts], dtype=np.int32), n_sites)
+    sites = result.sites
+    columns = {
+        "site": site_idx,
+        "kind": pa.array(
+            [sites[i].kind for i in site_idx], type=pa.dictionary(pa.int8(), pa.string())
+        ),
+        "layer": np.array([sites[i].layer for i in site_idx], dtype=np.int16),
+        "head": np.array(
+            [-1 if sites[i].head is None else sites[i].head for i in site_idx], dtype=np.int16
+        ),
+        "feature": np.array(
+            [-1 if sites[i].site.feature is None else sites[i].site.feature for i in site_idx],
+            dtype=np.int32,
+        ),
+        "position": pa.array(
+            [sites[i].position_key() for i in site_idx], type=pa.dictionary(pa.int16(), pa.string())
+        ),
+        "variant": pa.array(
+            [sites[i].variant_key or "" for i in site_idx],
+            type=pa.dictionary(pa.int16(), pa.string()),
+        ),
+        "prompt": prompt_idx,
+        "patched_logit_diff": result.patched_ld.ravel(),
+        "patched_answer_prob": result.patched_prob.ravel(),
+        "receiver_logit_diff": np.tile(result.receiver_ld, n_sites),
+        "reference_logit_diff": np.tile(result.reference_ld, n_sites),
+        "receiver_answer_prob": np.tile(result.receiver_prob, n_sites),
+        "delta": stats.delta.ravel(),
+        "effect": stats.effect.ravel(),
+    }
+    return pa.table(columns)
+
+
+def write_results(path: Path, table: pa.Table) -> None:
+    with atomic_output(path) as fh:
+        pq.write_table(table, fh, compression="zstd", write_statistics=False)
+
+
+def read_site_rows(path: Path, site: int) -> dict[str, list[Any]]:
+    table = pq.read_table(path, filters=[("site", "=", site)])
+    table = table.sort_by("prompt")
+    return {name: table.column(name).to_pylist() for name in table.column_names}