open-interpretability 0.0.20__py3-none-any.whl

openinterp/harness/scan.py

@@ -0,0 +1,462 @@
+"""One pass over a corpus that gathers everything an audit needs about a neuron pole.
+
+For every non-padding token, it records three things:
+
+- the pole strength, for the top examples and histogram;
+- the whole layer's activations, for partner neuron poles (co-activation lift, Spearman,
+  Pearson);
+- the token id and position, for the confound profile.
+
+Documents are scanned whole by default. With ``window_tokens``, each is cut into
+overlapping windows that start every ``stride`` tokens. So a long document is read in
+full without one huge forward pass. Every token still counts once in the statistics. It
+counts in the first window after the first that reaches it. There it has at least
+``window_tokens - stride`` tokens of context (the strided-perplexity convention).
+A unit with a first-token floor (protocol 0.8) is above its line at the first token,
+whatever the token. So its scan leaves first tokens out of its scale, histogram and
+top-example ranking (:func:`settle`). That is decided after the pass, so every scan
+keeps whether each counted token is its piece's first.
+Top examples are windows, decoded back to text so a record re-measures them as the
+scan did, and no two overlap.
+
+Each forward pass yields a :class:`Chunk`: the whole layer's activations on every real
+token of its pieces. Given a :class:`~openinterp.harness.layer_cache.LayerCache`, a scan
+stores its chunks, and a later scan of any neuron in that layer replays them.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Iterator, Sequence
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING
+
+import torch
+
+from openinterp.core.schemas import (
+    ActivatingExample,
+    ConfoundProfile,
+    CorrelatedNeuron,
+    Pole,
+    ReviewPolicy,
+)
+from openinterp.harness.annotations import LabelPass, Tagger, TokenLabels
+from openinterp.harness.confounds import ConfoundAccumulator
+from openinterp.harness.covariance import PartnerSearch
+from openinterp.harness.hooks import Inputs, NeuronHookRunner, pole_strength
+from openinterp.harness.scan_backend import forward_chunks, rescore_scan, with_ane_fallback
+
+if TYPE_CHECKING:
+    from openinterp.harness.layer_cache import LayerCache
+
+
+# Padded tokens per forward pass (pieces × the longest). Part of a cached scan's key, as
+# it sets padding. On Qwen2.5-0.5B over corpus.jsonl, the model reads 27% fewer tokens
+# (padding included) than with 8 pieces a pass. 2048 is near the fastest budget on
+# Apple's GPU.
+BATCH_TOKENS = 2048
+
+
+@dataclass
+class ScanResult:
+    """What one scan found: top examples, every counted strength, partners and confounds."""
+
+    examples: list[ActivatingExample]
+    values: list[float]  # every scanned token's pole strength, each token once
+    correlated: list[CorrelatedNeuron]
+    confounds: ConfoundProfile
+    windows: int = 0  # forward-pass inputs: documents, or windows of them
+    opposite: OppositeScan | None = None  # the neuron's other pole, from the same pass
+    # Protocol 0.8: first tokens are left out of examples and values (see ``settle``).
+    first_token_floor: bool = False
+
+
+@dataclass
+class OppositeScan:
+    """The neuron's other pole, from the same pass: its top examples and strengths."""
+
+    pole: Pole
+    examples: list[ActivatingExample]
+    values: list[float]
+    first_token_floor: bool = False
+
+
+Strengths = Callable[[Inputs], torch.Tensor]
+
+
+@torch.no_grad()
+def first_position_share(
+    runner: NeuronHookRunner, strengths: Strengths, line: float, batch_size: int = 8192
+) -> float:
+    """Return the share of the vocabulary on which ``strengths`` is at or above ``line``.
+
+    This is at a text's first position, with special tokens left out. Attention there
+    sees only the token itself. So this is a property of the weights, the same for every
+    corpus.
+    """
+    from openinterp.harness import first_position
+
+    if reader := getattr(strengths, "first_position", None):  # a neuron's: computed once
+        values = reader()
+    else:
+        values = first_position.strengths_of(runner, strengths, batch_size)
+    return float((values >= line).sum()) / max(len(values), 1)
+
+
+def has_first_token_floor(
+    runner: NeuronHookRunner, strengths: Strengths, scale: float, policy: ReviewPolicy
+) -> bool:
+    """Return whether the unit has a first-token floor (protocol 0.8).
+
+    It has one when it is active at the first position on at least
+    ``policy.max_active_share`` of the vocabulary. Its line comes from ``scale``.
+    """
+    if scale <= 0:
+        return False
+    line = policy.activation_threshold * scale
+    return first_position_share(runner, strengths, line) >= policy.max_active_share
+
+
+def settle(
+    runner: NeuronHookRunner,
+    strengths: Strengths,
+    values: Sequence[float],
+    first: Sequence[bool],
+    top: TopExamples,
+    k: int,
+    labels: TokenLabels | None = None,
+) -> tuple[list[ActivatingExample], list[float], bool]:
+    """Return a scan's top examples, counted values, and whether it has a first-token floor.
+
+    The floor is judged against the scale that the unit would have without first tokens.
+    That is the top example's peak after its first token. It is the scale that the unit
+    records if it has the floor. With the floor, first tokens are left out of the
+    examples' ranking and of the values. Each example carries its peak token's ``labels``.
+    """
+    best = top.top(runner.tokenizer, 1, skip_first=True)
+    scale = best[0].max_activation if best else 0.0
+    floor = has_first_token_floor(runner, strengths, scale, ReviewPolicy())
+    kept = [v for v, f in zip(values, first, strict=True) if not f] if floor else list(values)
+    return top.top(runner.tokenizer, k, skip_first=floor, labels=labels), kept, floor
+
+
+@dataclass(frozen=True)
+class Piece:
+    """What one forward pass reads: a whole document, or tokens [start, end) of it."""
+
+    doc: int
+    start: int
+    end: int
+    text: str
+    fresh: int | None = None  # trailing tokens new to the scan; None: every token is
+
+    def overlaps(self, other: Piece) -> bool:
+        """True when the two pieces share tokens of one document."""
+        return self.doc == other.doc and self.start < other.end and other.start < self.end
+
+
+def default_stride(window_tokens: int) -> int:
+    """Half a window, so every counted token has at least half a window of context."""
+    return max(window_tokens // 2, 1)
+
+
+def cut_pieces(
+    tokenizer, texts: Sequence[str], window_tokens: int | None, stride: int | None
+) -> Iterator[Piece]:
+    """Whole documents, or each document's windows (see the module docstring)."""
+    if window_tokens is None:
+        yield from (Piece(doc, 0, 1, text) for doc, text in enumerate(texts))
+        return
+    stride = stride or default_stride(window_tokens)
+    if not 0 < stride <= window_tokens:
+        raise ValueError(f"stride must be in 1..{window_tokens}, got {stride}")
+    for doc, text in enumerate(texts):
+        ids = tokenizer.encode(text, add_special_tokens=False)
+        start = prev_end = 0
+        while start < len(ids):
+            end = min(start + window_tokens, len(ids))
+            fresh = None if start == 0 else end - prev_end
+            yield Piece(doc, start, end, tokenizer.decode(ids[start:end]), fresh)
+            if end == len(ids):
+                break
+            start, prev_end = start + stride, end
+
+
+@dataclass
+class Batch:
+    """One forward pass of pieces, with which of its tokens the scan counts."""
+
+    pieces: list[Piece]
+    encoded: dict[str, torch.Tensor]
+    ids: torch.Tensor  # [b, seq], on the CPU
+    mask: torch.Tensor  # real (non-padding) tokens
+    counted: torch.Tensor  # real tokens new to the scan
+    positions: torch.Tensor
+
+    @property
+    def first(self) -> torch.Tensor:
+        """Each piece's first real token."""
+        return self.mask & (self.mask.cumsum(1) == 1)
+
+
+def plan_batches(lengths: Sequence[int], batch_tokens: int) -> list[int]:
+    """Return how many consecutive pieces each forward pass takes.
+
+    That is as many as fit in ``batch_tokens`` once padded to the longest, and at least one.
+    """
+    sizes: list[int] = []
+    size = longest = 0
+    for n in lengths:
+        if size and max(longest, n) * (size + 1) > batch_tokens:
+            sizes.append(size)
+            size = longest = 0
+        size, longest = size + 1, max(longest, n)
+    return [*sizes, size] if size else sizes
+
+
+def batches(
+    runner: NeuronHookRunner, pieces: Sequence[Piece], batch_tokens: int = BATCH_TOKENS
+) -> Iterator[Batch]:
+    """``pieces`` in order, in forward passes of at most ``batch_tokens`` padded tokens."""
+    if not pieces:
+        return
+    every = runner.tokenizer([p.text for p in pieces], return_tensors="pt", padding=True)
+    lengths = every.get("attention_mask", torch.ones_like(every["input_ids"])).sum(1).tolist()
+    start = 0
+    for size in plan_batches(lengths, batch_tokens):
+        batch = list(pieces[start : start + size])
+        start += size
+        encoded = runner.encode([p.text for p in batch])
+        ids = encoded["input_ids"].cpu()
+        mask = encoded.get("attention_mask", torch.ones_like(encoded["input_ids"])).bool().cpu()
+        # Count only each piece's fresh tail: real tokens at or after it ≤ fresh.
+        fresh = torch.tensor([p.fresh or ids.shape[1] for p in batch]).unsqueeze(1)
+        counted = mask & (mask.flip(1).cumsum(1).flip(1) <= fresh)
+        positions = torch.arange(ids.shape[1]).expand_as(ids)
+        yield Batch(batch, encoded, ids, mask, counted, positions)
+
+
+@dataclass
+class Chunk:
+    """One forward pass, flattened to its pieces' real tokens in scan order."""
+
+    pieces: list[Piece]
+    lengths: list[int]  # real tokens of each piece
+    ids: torch.Tensor  # [n]
+    positions: torch.Tensor  # [n]
+    counted: torch.Tensor  # [n] bool: tokens new to the scan
+    acts: torch.Tensor  # [n, width] the whole layer, in the model's dtype, on the CPU
+
+    @property
+    def first(self) -> torch.Tensor:
+        """``[n]`` bool: each piece's first token (pieces are padded on the right)."""
+        return self.positions == 0
+
+    def rows(self) -> Iterator[tuple[Piece, slice]]:
+        """Each piece and its rows."""
+        start = 0
+        for piece, n in zip(self.pieces, self.lengths, strict=True):
+            yield piece, slice(start, start + n)
+            start += n
+
+
+def layer_chunks(
+    runner: NeuronHookRunner, pieces: Sequence[Piece], layer: int, batch_tokens: int
+) -> Iterator[Chunk]:
+    """Run the model over ``pieces`` and yield ``layer``'s activations batch by batch."""
+    for batch in batches(runner, pieces, batch_tokens):
+        cache = runner.layer_activations(batch.encoded, [layer])
+        yield Chunk(
+            pieces=batch.pieces,
+            lengths=batch.mask.sum(1).tolist(),
+            ids=batch.ids[batch.mask],
+            positions=batch.positions[batch.mask],
+            counted=batch.counted[batch.mask],
+            acts=cache[layer].cpu()[batch.mask],
+        )
+
+
+def scan_chunks(
+    runner: NeuronHookRunner,
+    pieces: Sequence[Piece],
+    layer: int,
+    batch_tokens: int = BATCH_TOKENS,
+    cache: LayerCache | None = None,
+) -> Iterator[Chunk]:
+    """Yield ``layer``'s chunks, replayed from ``cache`` if it holds them.
+
+    Otherwise they come from the model, and are stored in ``cache``.
+    """
+    chunks = forward_chunks(runner, pieces, layer, batch_tokens)  # on the ANE if it can
+    return chunks if cache is None else cache.through(runner, layer, pieces, batch_tokens, chunks)
+
+
+class TopExamples:
+    """Collects every piece's strengths and keeps the ``k`` strongest that don't overlap."""
+
+    def __init__(self) -> None:
+        self.candidates: list[tuple[Piece, list[int], list[float]]] = []
+
+    def update(self, batch: Batch, strengths: torch.Tensor) -> None:
+        """Add a batch's pieces as candidates, with their real tokens' strengths."""
+        for row, piece in enumerate(batch.pieces):
+            mask = batch.mask[row]
+            self.add(piece, batch.ids[row][mask].tolist(), strengths[row][mask].tolist())
+
+    def update_chunk(self, chunk: Chunk, strengths: torch.Tensor) -> None:
+        """Add a chunk's pieces as candidates, with their strengths."""
+        for piece, rows in chunk.rows():
+            self.add(piece, chunk.ids[rows].tolist(), strengths[rows].tolist())
+
+    def add(self, piece: Piece, ids: list[int], values: list[float]) -> None:
+        """Add one piece as a candidate, unless it has no strengths."""
+        if values:
+            self.candidates.append((piece, ids, values))
+
+    def top(
+        self, tokenizer, k: int, skip_first: bool = False, labels: TokenLabels | None = None
+    ) -> list[ActivatingExample]:
+        """Return the ``k`` strongest pieces by their peak.
+
+        With ``skip_first``, the peak is after the first token. Each has its peak token's
+        ``labels``.
+        """
+        start = int(skip_first)
+        ranked = [(max(v[start:]), p, i, v) for p, i, v in self.candidates if len(v) > start]
+        examples: list[ActivatingExample] = []
+        chosen: list[Piece] = []
+        for peak_value, piece, row_ids, row_vals in sorted(ranked, key=lambda c: -c[0]):
+            if len(examples) == k:
+                break
+            if any(piece.overlaps(c) for c in chosen):
+                continue
+            chosen.append(piece)
+            peak = row_vals.index(peak_value, start)
+            examples.append(
+                ActivatingExample(
+                    text=piece.text,
+                    tokens=[tokenizer.decode([i]) for i in row_ids],
+                    activations=[round(v, 5) for v in row_vals],
+                    max_activation=peak_value,
+                    max_token_index=peak,
+                    peak_labels=labels.at(piece, row_ids, peak) if labels else None,
+                )
+            )
+        return examples
+
+
+def scan_neuron(
+    runner: NeuronHookRunner,
+    texts: Sequence[str],
+    layer: int,
+    neuron_idx: int,
+    pole: Pole,
+    *,
+    k: int = 20,
+    batch_tokens: int = BATCH_TOKENS,
+    n_correlated: int = 10,
+    min_count: int = 3,
+    window_tokens: int | None = None,
+    stride: int | None = None,
+    cache: LayerCache | None = None,
+    opposite_k: int = 5,
+    taggers: Sequence[Tagger] = (),
+) -> ScanResult:
+    """Scan ``texts`` for one neuron pole, and for its other pole.
+
+    The other pole gets its top ``opposite_k`` examples and histogram. That is free,
+    because the pass reads the whole activation. With ``cache``, replay the layer's stored
+    activations if it holds this scan, and store them if not. Each of ``taggers`` labels
+    the counted tokens, and its confound joins the profile's.
+    """
+    if runner.tokenizer is None:
+        raise ValueError("scan_neuron requires a tokenizer")
+    other: Pole = "negative" if pole == "positive" else "positive"
+    pieces = list(cut_pieces(runner.tokenizer, texts, window_tokens, stride))
+
+    def run():  # fresh accumulators, so a pass redone on torch starts from nothing
+        labels = LabelPass(runner.tokenizer, texts, taggers)
+        poles = (pole, other)
+        return labels, _pole_pass(
+            runner, pieces, layer, neuron_idx, poles, batch_tokens, cache, labels
+        )
+
+    labels, scan = with_ane_fallback(runner, run)
+    tops = (scan.top, scan.other_top)
+    rescore_scan(
+        runner, pieces, layer, neuron_idx, (pole, other), tops, (k, opposite_k), batch_tokens
+    )
+    read, read_other = (_reader(runner, layer, neuron_idx, p) for p in (pole, other))
+    tokens = labels.tokens
+    examples, values, floor = settle(runner, read, scan.values, scan.first, scan.top, k, tokens)
+    other_examples, other_values, other_floor = settle(
+        runner, read_other, scan.other_values, scan.first, scan.other_top, opposite_k, tokens
+    )
+    # "Active" for co-activation uses the protocol's own definition for the target.
+    active = ReviewPolicy().activation_threshold * max(values, default=0.0)
+    confounds = labels.finalize(scan.confounds.finalize(runner.tokenizer, min_count), values)
+    return ScanResult(
+        examples=examples,
+        values=values,
+        correlated=scan.partners.partners(neuron_idx, active, k=n_correlated) if active > 0 else [],
+        confounds=confounds,
+        windows=len(pieces),
+        opposite=OppositeScan(other, other_examples, other_values, other_floor),
+        first_token_floor=floor,
+    )
+
+
+@dataclass
+class _PolePass:
+    """What one pass over the pieces gathers for a neuron's pole and its other pole."""
+
+    partners: PartnerSearch
+    confounds: ConfoundAccumulator = field(default_factory=ConfoundAccumulator)
+    top: TopExamples = field(default_factory=TopExamples)
+    other_top: TopExamples = field(default_factory=TopExamples)
+    values: list[float] = field(default_factory=list)
+    other_values: list[float] = field(default_factory=list)
+    first: list[bool] = field(default_factory=list)
+
+
+def _pole_pass(runner, pieces, layer: int, neuron_idx: int, poles, batch_tokens, cache, labels):
+    pole, other = poles
+    scan = _PolePass(PartnerSearch(runner.mlp_width(layer)))
+    for chunk in scan_chunks(runner, pieces, layer, batch_tokens, cache):
+        acts = chunk.acts.float()  # [n, d_mlp]
+        strengths = pole_strength(acts[:, neuron_idx], pole)
+        counted = chunk.counted
+        scan.partners.update(strengths[counted], acts[counted])
+        scan.confounds.update(chunk.ids[counted], chunk.positions[counted], strengths[counted])
+        labels.update_chunk(chunk, strengths)
+        scan.values.extend(strengths[counted].tolist())
+        scan.first.extend(chunk.first[counted].tolist())
+        scan.top.update_chunk(chunk, strengths)
+        other_strengths = pole_strength(acts[:, neuron_idx], other)
+        scan.other_values.extend(other_strengths[counted].tolist())
+        scan.other_top.update_chunk(chunk, other_strengths)
+    return scan
+
+
+def _reader(runner, layer: int, neuron_idx: int, pole: Pole) -> Strengths:
+    from openinterp.harness.first_position import NeuronReader
+
+    return NeuronReader(runner, layer, neuron_idx, pole)
+
+
+def top_activating_examples(
+    runner: NeuronHookRunner,
+    texts: Sequence[str],
+    layer: int,
+    neuron_idx: int,
+    pole: Pole,
+    *,
+    k: int = 20,
+    batch_tokens: int = BATCH_TOKENS,
+) -> tuple[list[ActivatingExample], list[float]]:
+    """Scan ``texts``; return the ``k`` strongest snippets and every token's pole strength.
+
+    A thin wrapper over :func:`scan_neuron`, which also gathers correlated neurons and
+    the confound profile.
+    """
+    result = scan_neuron(runner, texts, layer, neuron_idx, pole, k=k, batch_tokens=batch_tokens)
+    return result.examples, result.values

openinterp/harness/scan_backend.py

@@ -0,0 +1,201 @@
+"""Where a scan's forward pass runs, and the exact measurement of its top candidates.
+
+A runner can carry a ``scan_backend``, an :class:`~openinterp.harness.ane.backend.AneScan`.
+It runs the scan's forward pass on Apple's Neural Engine (ANE). The device planner
+(``harness/devices.py``) adds it when the GPU is busy. The ANE computes in fp16, about 1%
+from fp32. So after the pass, :func:`rescore_scan` measures each pole's strongest
+candidates again on the runner's own device:
+
+- the :data:`RESCORE_TOP` strongest pieces;
+- and any piece within :data:`RESCORE_MARGIN` of the 20th strongest, up to
+  :data:`RESCORE_CAP` pieces.
+
+On ``corpus.jsonl`` that rule found the exact top 20 for every neuron and both poles. The
+layers were GPT-2 L5, Qwen2.5-0.5B L12 and Qwen2.5-1.5B L8 (docs/analyses/ane-feasibility.md).
+Then it makes sure that every example the scan will keep was measured again. The
+numbers that ``open-interpretability reproduce`` checks are then the device's numbers. The
+histogram, partners and confounds keep the ANE's values, which correlate with fp32 at
+r ≥ 0.9998 for every neuron.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Iterator, Sequence
+from typing import TYPE_CHECKING
+
+import click
+import torch
+
+from openinterp.harness.hooks import pole_strength
+
+if TYPE_CHECKING:
+    from openinterp.core.schemas import Pole
+    from openinterp.harness.hooks import NeuronHookRunner
+    from openinterp.harness.scan import Chunk, Piece, TopExamples
+
+RESCORE_TOP = 60
+RESCORE_MARGIN = 0.06
+RESCORE_CAP = 120
+RESCORE_RANK = 20  # the margin is below this rank's peak
+RESCORE_BATCH = 16  # pieces per forward pass when measuring again
+RESCORE_ROUNDS = 5  # rounds that add examples the selection would keep unmeasured
+
+
+def ane_scan(runner: NeuronHookRunner, pieces: Sequence[Piece], layer: int, batch_tokens: int):
+    """The runner's ANE backend, made ready for this scan, if it can run it; else None.
+
+    If it can't, it says why on stderr, once for each reason.
+    """
+    backend = getattr(runner, "scan_backend", None)
+    if backend is None or not pieces:
+        return None
+    every = runner.tokenizer([p.text for p in pieces], return_tensors="pt", padding=True)
+    mask = every.get("attention_mask", torch.ones_like(every["input_ids"]))
+    reason = backend.prepare(mask.sum(1).tolist(), layer, batch_tokens)
+    if reason is None:
+        return backend
+    if reason not in backend.noted:
+        backend.noted.add(reason)
+        click.echo(f"This scan runs on {runner.device.type}, not the ANE: {reason}.", err=True)
+    return None
+
+
+class AneError(RuntimeError):
+    """The ANE failed during a scan. The scan is then done again on torch."""
+
+
+def with_ane_fallback(runner: NeuronHookRunner, run: Callable):
+    """``run()``; if the ANE fails during it, stop using the ANE and ``run()`` again.
+
+    The second run starts from nothing, so no partial result of the ANE is kept.
+    """
+    try:
+        return run()
+    except AneError as exc:
+        backend = runner.scan_backend
+        backend.broken = f"it failed during a scan ({exc.__cause__ or exc})"
+        click.echo(
+            f"The ANE failed, so this scan runs again on {runner.device.type}: {exc}", err=True
+        )
+        return run()
+
+
+def scan_device(runner: NeuronHookRunner, pieces: Sequence[Piece], layer, batch_tokens) -> str:
+    """The device a scan of ``pieces`` runs on: the backend's name, or the runner's device."""
+    backend = ane_scan(runner, pieces, layer, batch_tokens)
+    return backend.name if backend is not None else runner.device.type
+
+
+def forward_chunks(
+    runner: NeuronHookRunner, pieces: Sequence[Piece], layer: int, batch_tokens: int
+) -> Iterator[Chunk]:
+    """``layer``'s chunks over ``pieces``, from the ANE if it can run them, else torch."""
+    from openinterp.harness.scan import layer_chunks
+
+    backend = ane_scan(runner, pieces, layer, batch_tokens)
+    if backend is None:
+        return layer_chunks(runner, pieces, layer, batch_tokens)
+    backend.used = True
+    return backend.chunks(runner, pieces, layer, batch_tokens)
+
+
+def rescore_scan(
+    runner, pieces, layer: int, neuron_idx: int, poles, tops, ks, batch_tokens
+) -> None:
+    """Measure again the top candidates of each pole, if the scan's values are the ANE's.
+
+    ``poles``, ``tops`` and ``ks`` are each pole, its :class:`TopExamples` and the number
+    of examples it keeps.
+    """
+    if ane_scan(runner, pieces, layer, batch_tokens) is None:
+        return
+    for pole, top, k in zip(poles, tops, ks, strict=True):
+        rescore(top, _measure(runner, layer, neuron_idx, pole), k)
+
+
+def _measure(runner, layer: int, neuron_idx: int, pole: Pole) -> Callable:
+    """A function that measures the pole's strengths on texts, :data:`RESCORE_BATCH` a pass.
+
+    On a model that streams, the passes run layer-major (``harness/layer_major.py``), so
+    each layer is read from disk once for all of them.
+    """
+    from openinterp.harness.layer_major import map_layer_major, prompt_tokens
+
+    def one(texts: list[str]) -> list[list[float]]:
+        encoded = runner.encode(texts)
+        mask = encoded.get("attention_mask", torch.ones_like(encoded["input_ids"])).bool()
+        with torch.no_grad():
+            acts = runner.layer_activations(encoded, [layer])[layer][..., neuron_idx]
+        values = pole_strength(acts.float(), pole).cpu()
+        return [values[r][mask[r].cpu()].tolist() for r in range(len(texts))]
+
+    def measure(texts: list[str]) -> list[list[float]]:
+        parts = [texts[i : i + RESCORE_BATCH] for i in range(0, len(texts), RESCORE_BATCH)]
+        tokens = RESCORE_BATCH * prompt_tokens(runner, texts)
+        done = map_layer_major(runner, one, parts, tokens=tokens, logits=False)
+        return [values for part in done for values in part]
+
+    return measure
+
+
+def rescore(top: TopExamples, measure: Callable[[list[str]], list[list[float]]], k: int) -> int:
+    """Replace the values of ``top``'s strongest candidates with ``measure``'s.
+
+    Return how many candidates it measured again.
+    """
+    done: set[int] = set()
+    todo = chosen(top.candidates)
+    for _ in range(RESCORE_ROUNDS):
+        todo -= done
+        if not todo:
+            break
+        order = sorted(todo)
+        new = measure([top.candidates[i][0].text for i in order])  # all of them at once
+        for i, values in zip(order, new, strict=True):
+            piece, ids, old = top.candidates[i]
+            if len(values) != len(old):
+                raise ValueError(f"measured {len(values)} tokens of a {len(old)}-token piece")
+            top.candidates[i] = (piece, ids, values)
+        done |= todo
+        todo = kept(top.candidates, k)
+    return len(done)
+
+
+def chosen(candidates: list) -> set[int]:
+    """Indices of the candidates to measure again, by their peak with and without token 0."""
+    picked: set[int] = set()
+    for start in (0, 1):
+        peaks = torch.tensor(
+            [max(v[start:]) if len(v) > start else -torch.inf for *_, v in candidates]
+        )
+        order = peaks.argsort(descending=True)
+        ranked = peaks[order]
+        if not len(ranked):
+            continue
+        nth = ranked[min(RESCORE_RANK, len(ranked)) - 1]
+        line = nth - RESCORE_MARGIN * nth.abs() - 1e-3 * ranked[0].abs()
+        for rank, index in enumerate(order.tolist()):
+            if rank < RESCORE_TOP or (ranked[rank] >= line and rank < RESCORE_CAP):
+                picked.add(index)
+    return picked
+
+
+def kept(candidates: list, k: int) -> set[int]:
+    """Indices of the candidates that the top ``k`` would keep, with and without token 0.
+
+    As ``TopExamples.top``, it skips a piece that overlaps one it kept.
+    """
+    keep: set[int] = set()
+    for start in (0, 1):
+        ranked = sorted(
+            (i for i, (*_, v) in enumerate(candidates) if len(v) > start),
+            key=lambda i: -max(candidates[i][2][start:]),
+        )
+        taken: list[int] = []
+        for i in ranked:
+            if len(taken) == k:
+                break
+            if not any(candidates[i][0].overlaps(candidates[j][0]) for j in taken):
+                taken.append(i)
+        keep.update(taken)
+    return keep