feat(magic): per-token per-query MAGIC, and score-format cleanups

bergson/config/config_io.py

@@ -120,6 +120,31 @@ def read_config(path: str | Path) -> dict[str, Any]:
     return config
 
 
+def read_first_step_config(path: str | Path) -> dict[str, Any] | None:
+    """The first step's configuration from a run's ``config.yaml``.
+
+    ``path`` may be the run directory or an artifact inside it (``scores.pt``),
+    so callers holding an output path need not resolve it. Returns ``None``
+    when there is no readable config, which is the signal that a caller must
+    fall back to inspecting the artifact itself.
+    """
+    p = Path(path)
+    cfg_path = (p if p.is_dir() else p.parent) / CONFIG_FILENAME
+    if not cfg_path.is_file():
+        return None
+
+    try:
+        steps = read_config(cfg_path)["steps"]
+    except (ValueError, yaml.YAMLError):
+        return None
+
+    for step in steps:
+        for step_cfg in step.values():
+            if isinstance(step_cfg, dict):
+                return step_cfg
+    return None
+
+
 T = TypeVar("T", bound="FromDict")
 
 

bergson/data.py

@@ -640,40 +640,13 @@ def to_grid(self) -> torch.Tensor:
         return scores[..., 0] if self.num_scores == 1 else scores
 
 
-class ArrayScores:
-    """Dense ``[num_items, num_scores]`` score matrix with the same interface
-    as :class:`Scores`, for scores saved as a plain ``.npy`` (e.g. an ad hoc
-    array from outside the standard scoring pipeline)."""
-
-    def __init__(self, arr: np.ndarray):
-        if arr.ndim == 1:
-            arr = arr[:, None]
-        self.arr = arr
-        self.num_scores = arr.shape[1]
-
-    def __len__(self) -> int:
-        return self.arr.shape[0]
-
-    def __getitem__(self, key: Any) -> NDArray:
-        return self.arr[key]
-
-    def get(self, key: Any, score_idx: int = 0) -> NDArray:
-        return self.arr[key, score_idx]
-
-    def is_written(self) -> bool:
-        return True
-
-
-def load_scores(path: Path) -> Scores | ArrayScores:
+def load_scores(path: Path) -> Scores:
     """Load a score store written by the standard scoring pipeline.
 
-    A plain ``.npy`` array loads as :class:`ArrayScores`. Any other score
-    directory loads its ``scores.bin`` as :class:`Scores`, with ``offsets``
-    set when ``info["attribute_tokens"]`` marks it as a per-token store.
+    A score directory loads its ``scores.bin`` as :class:`Scores`, with
+    ``offsets`` set when ``info["attribute_tokens"]`` marks it as a per-token
+    store.
     """
-    if path.suffix == ".npy":
-        return ArrayScores(np.load(path))
-
     info_path = path / "info.json"
     with open(info_path, "r") as f:
         info = json.load(f)

bergson/magic/cli.py

@@ -28,7 +28,7 @@
 )
 
 from ..config.config import TrainingConfig, ValidationConfig
-from ..config.config_io import save_run_config
+from ..config.config_io import read_first_step_config, save_run_config
 from ..distributed import launch_distributed_run
 from ..utils.load_from_optimizer import (
     save_second_moments_as_optimizer_pt,
@@ -130,9 +130,8 @@ def compute_per_query_magic_scores(
     yields a single aggregate-query score. This scores each query separately:
     for each query document it takes that document's gradient at the final model
     as the backward cotangent and runs ``Trainer.backward`` over the saved
-    trajectory, producing a ``[num_train_docs, num_query_docs]`` matrix (the
-    layout ``validate_scores`` consumes: rows are leave-out docs, columns
-    queries).
+    trajectory, producing ``[num_train_docs, num_query_docs]`` — or
+    ``[num_train_docs, seq_len, num_query_docs]`` when attributing tokens.
 
     The backward is linear in the cotangent, so this is exact; the forward runs
     once and every query reuses its checkpoints (``cleanup=False``). Per-query
@@ -217,7 +216,7 @@ def compute_per_query_magic_scores(
 
         s = bwd_state.weight_grads.detach().cpu()
         if pad_count:
-            s = s[:-weight_pad_count]
+            s = s[:-weight_pad_count] if s.ndim == 1 else s[:-pad_count]
         if main:
             torch.save(s, qpath)
         per_query.append(s)
@@ -241,9 +240,7 @@ def compute_per_query_magic_scores(
     fwd_state.copy_(restored)
     del restored
 
-    # [num_train_docs, num_query_docs] — the layout validate_scores expects
-    # (rows are leave-out docs; columns are queries).
-    return torch.stack(per_query, dim=1)
+    return torch.stack(per_query, dim=-1)
 
 
 def scores_are_per_token(score_path: str) -> bool:
@@ -253,10 +250,14 @@ def scores_are_per_token(score_path: str) -> bool:
             return False
         with open(info_path) as f:
             return bool(json.load(f).get("attribute_tokens", False))
-    if score_path.endswith(".npy"):
-        return False
+    step_cfg = read_first_step_config(score_path)
+    if step_cfg is not None:
+        return bool(step_cfg.get("attribute_tokens") or step_cfg.get("per_token"))
+
     scores = torch.load(score_path, map_location="cpu")
-    return isinstance(scores, torch.Tensor) and scores.ndim == 2 and scores.shape[1] > 1
+    return isinstance(scores, torch.Tensor) and (
+        scores.ndim == 3 or (scores.ndim == 2 and scores.shape[1] > 1)
+    )
 
 
 def attach_doc_ids_if_missing(dataset: Dataset) -> Dataset:
@@ -487,9 +488,8 @@ def worker(
 
     multi_query = False
     if not score_path and run_cfg.query_method == "none":
-        # Per-query MAGIC: one backward per query, sharing the forward. Yields a
-        # [num_query_docs, num_train_docs] score matrix (multi_query), the unit
-        # for a per-query LDS.
+        # Per-query MAGIC: one backward per query, sharing the forward. Yields
+        # the unit for a per-query LDS.
         if not isinstance(run_cfg, MagicConfig):
             raise RuntimeError("run_cfg must be a MagicConfig to compute scores")
         assert query_dataset is not None
@@ -569,7 +569,7 @@ def worker(
             score_path = os.path.join(run_cfg.run_path, "scores.pt")
             torch.save(scores, score_path)
             print(f"Saved attribution scores to {score_path}")
-    elif os.path.isdir(score_path) or score_path.endswith(".npy"):
+    elif os.path.isdir(score_path):
         scores, multi_query = load_attribution_scores(score_path)
     else:
         scores = torch.load(score_path, map_location="cpu")

bergson/score/score.py

@@ -231,7 +231,7 @@ def create_scorer(
     else:
         num_scores = num_queries_total if num_queries_total is not None else num_queries
     if attribute_tokens:
-        writer = MemmapTokenScoreWriter(
+        writer = MemmapTokenScoreWriter.from_dataset(
             path,
             data,
             num_scores,

bergson/score/score_writer.py

@@ -156,22 +156,22 @@ class MemmapTokenScoreWriter(ScoreWriter):
     def __init__(
         self,
         path: Path,
-        data: Dataset,
+        num_token_grads: np.ndarray,
         num_scores: int,
         *,
         dtype: torch.dtype = torch.float32,
         flush_interval: int = 64,
+        overwrite: bool = False,
     ):
         self.path = path
         self.num_scores = num_scores
         self.dtype = dtype
         self.flush_interval = flush_interval
         self.num_batches_since_flush = 0
 
-        num_token_grads = compute_num_token_grads(data)
-        num_items = len(data)
+        num_items = len(num_token_grads)
         self.num_token_grads = num_token_grads
-        self.offsets = np.zeros(len(num_token_grads) + 1, dtype=np.int64)
+        self.offsets = np.zeros(num_items + 1, dtype=np.int64)
         np.cumsum(num_token_grads, out=self.offsets[1:])
         total_tokens = int(self.offsets[-1])
 
@@ -181,7 +181,7 @@ def __init__(
         struct_dtype, struct_dtype_json = _score_struct_dtype(num_scores, np_dtype)
 
         rank = dist.get_rank() if dist.is_initialized() else 0
-        if rank == 0 and not scores_file_path.exists():
+        if rank == 0 and (overwrite or not scores_file_path.exists()):
             print(f"Creating new token scores file: {scores_file_path}")
 
             self.scores = np.memmap(
@@ -216,6 +216,11 @@ def __init__(
             shape=(total_tokens,),
         )
 
+    @classmethod
+    def from_dataset(cls, path: Path, data: Dataset, num_scores: int, **kwargs):
+        """For callers holding the dataset the scores were computed over."""
+        return cls(path, compute_num_token_grads(data), num_scores, **kwargs)
+
     def __call__(self, indices: list[int], scores: torch.Tensor, query_offset: int = 0):
         # scores: [total_valid_in_batch, num_scores]
         scores = scores.to(dtype=self.dtype)
@@ -248,48 +253,24 @@ def save_token_scores(
     *,
     dtype: torch.dtype = torch.float32,
 ) -> None:
-    """One-shot equivalent of :class:`MemmapTokenScoreWriter`'s on-disk
-    layout, for callers that already hold the full flat ``(total_tokens,
-    num_scores)`` array in memory (e.g. summed across upstream per-token
-    score dirs sharing the same ``offsets``), rather than streaming batches
-    through ``__call__``.
+    """One-shot equivalent of :class:`MemmapTokenScoreWriter`, for callers that
+    already hold the full flat ``(total_tokens, num_scores)`` array in memory
+    (e.g. summed across upstream per-token score dirs sharing the same
+    ``offsets``), rather than streaming batches through ``__call__``.
     """
     if scores.ndim == 1:
         scores = scores[:, None]
-    total_tokens, num_scores = scores.shape
-    num_items = len(offsets) - 1
-
-    path.mkdir(parents=True, exist_ok=True)
-    np_dtype = convert_dtype_to_np(dtype)
-    struct_dtype, struct_dtype_json = _score_struct_dtype(num_scores, np_dtype)
-
-    mmap = np.memmap(
-        path / "scores.bin",
-        dtype=np.dtype(struct_dtype),  # type: ignore
-        mode="w+",
-        shape=(total_tokens,),
+    _, num_scores = scores.shape
+    num_token_grads = np.diff(offsets).astype(np.int64)
+
+    writer = MemmapTokenScoreWriter(
+        path, num_token_grads, num_scores, dtype=dtype, overwrite=True
     )
-    # numpy_to_tensor handles bf16, which torch.from_numpy cannot.
-    scores_np = tensor_to_numpy(numpy_to_tensor(np.ascontiguousarray(scores)).to(dtype))
-    for i in range(num_scores):
-        mmap[f"score_{i}"] = scores_np[:, i]
-        mmap[f"written_{i}"] = True
-    mmap.flush()
-
-    np.save(path / "offsets.npy", offsets)
-
-    with (path / "info.json").open("w") as f:
-        json.dump(
-            {
-                "attribute_tokens": True,
-                "num_items": num_items,
-                "num_rows": total_tokens,
-                "num_scores": num_scores,
-                "dtype": struct_dtype_json,
-            },
-            f,
-            indent=2,
-        )
+    writer(
+        list(range(len(num_token_grads))),
+        numpy_to_tensor(np.ascontiguousarray(scores)),
+    )
+    writer.flush()
 
 
 class MemmapSequenceScoreWriter(ScoreWriter):

bergson/validate.py

@@ -18,7 +18,6 @@
 import torch
 import torch.distributed as dist
 import torch.nn.functional as F
-import yaml
 from peft import PeftModel
 from scipy.stats import pearsonr, spearmanr
 from tqdm import tqdm
@@ -31,7 +30,7 @@
 )
 
 from .config.config import ScoreConfig, ValidationConfig
-from .config.config_io import CONFIG_FILENAME, load_subconfig, save_run_config
+from .config.config_io import load_subconfig, read_first_step_config, save_run_config
 from .data import Scores, load_scores, pad_and_tensor
 from .magic.data_stream import DataStream, pad_dataset_to_batch_size
 from .magic.trainer import TrainerState, prepare_trainer
@@ -41,71 +40,47 @@
 
 
 def load_attribution_scores(score_path: str) -> tuple[torch.Tensor, bool]:
-    """Load attribution scores from a score directory, ``.npy``, or ``.pt`` file.
+    """Load attribution scores from a score directory or ``.pt`` file.
 
     Returns ``(scores, multi_query)``. Token score directories are per-token,
     with a query dimension when ``num_scores > 1`` (``[docs, seq_len,
-    queries]``); plain score directories and ``.npy`` arrays are per-document,
-    with one column per query. A 2-D ``.pt`` tensor is ambiguous — per-token
-    MAGIC scores are ``[docs, seq_len]`` and per-query MAGIC scores are
-    ``[docs, queries]`` — so the run config next to it decides: it is
-    per-query iff the run used ``query_method: none``.
+    queries]``); plain score directories are per-document, with one column per
+    query. A 3-D ``.pt`` tensor is per-token per-query ``[docs, seq_len,
+    queries]``. A 2-D ``.pt`` tensor is ambiguous — per-token MAGIC scores are
+    ``[docs, seq_len]`` and per-query MAGIC scores are ``[docs, queries]`` —
+    so the run config next to it decides: it is per-query iff the run used
+    ``query_method: none`` without attributing tokens.
 
     Score directories are negated when their ``score_cfg.higher_is_better`` is
-    set, aligning them with the loss-diff convention. ``.npy`` files carry no
-    ``score_cfg`` and are loaded as-is: they must already be in the loss-diff
-    convention.
+    set, aligning them with the loss-diff convention.
     """
-    if os.path.isdir(score_path) or score_path.endswith(".npy"):
+    if os.path.isdir(score_path):
         loaded = load_scores(Path(score_path))
-        score_cfg = (
-            load_subconfig(score_path, "score_cfg", ScoreConfig)
-            if os.path.isdir(score_path)
-            else None
-        )
+        score_cfg = load_subconfig(score_path, "score_cfg", ScoreConfig)
         negate = score_cfg is not None and score_cfg.higher_is_better
 
         if isinstance(loaded, Scores) and loaded.offsets is not None:
             scores = loaded.to_grid()
             if negate:
                 scores = -scores
-            return scores, scores.ndim == 3
+            return scores, loaded.num_scores > 1
 
         arr = np.asarray(loaded[:])
         # Copy: the slice is a read-only view onto the memmap.
         out_dtype = arr.dtype if np.issubdtype(arr.dtype, np.floating) else np.float32
         scores = torch.from_numpy(arr.astype(out_dtype, copy=True))
         if negate:
             scores = -scores
-        return scores, scores.ndim == 2 and scores.shape[1] > 1
+        return scores, loaded.num_scores > 1
 
     scores = torch.load(score_path, map_location="cpu")
-    return scores, _pt_scores_are_per_query(score_path, scores)
-
-
-def _pt_scores_are_per_query(score_path: str, scores: torch.Tensor) -> bool:
-    """Whether a 2-D ``.pt`` tensor is per-query ``[docs, queries]`` rather
-    than per-token ``[docs, seq_len]``. The shape alone cannot tell them
-    apart; the run config written next to ``scores.pt`` records which one
-    ``run_magic`` saved."""
-    if not (isinstance(scores, torch.Tensor) and scores.ndim == 2):
-        return False
-    cfg_path = Path(score_path).parent / CONFIG_FILENAME
-    if not cfg_path.is_file():
-        return False
-    with open(cfg_path) as f:
-        doc = yaml.safe_load(f)
-    if not isinstance(doc, dict):
-        return False
-    steps = doc.get("steps")
-    payload = (
-        next(iter(steps[0].values())) if isinstance(steps, list) and steps else doc
-    )
-    return (
-        isinstance(payload, dict)
-        and payload.get("query_method") == "none"
-        and not payload.get("per_token")
-    )
+    if not isinstance(scores, torch.Tensor) or scores.ndim not in (2, 3):
+        return scores, False
+
+    step_cfg = read_first_step_config(score_path)
+    if step_cfg is None:
+        return scores, scores.ndim == 3
+    return scores, step_cfg.get("query_method") == "none"
 
 
 def bank_loss_cache_key(