Write scores as a score directory

bergson/config/config_io.py

@@ -120,31 +120,6 @@ 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

@@ -13,6 +13,7 @@
 import pyarrow as pa
 import torch
 import torch.distributed as dist
+import yaml
 from datasets import (
     Dataset,
     DatasetDict,
@@ -25,7 +26,8 @@
 from numpy.typing import DTypeLike, NDArray
 from transformers import PreTrainedTokenizerFast, logging
 
-from .config.config import DataConfig
+from .config.config import DataConfig, ScoreConfig
+from .config.config_io import CONFIG_FILENAME, load_subconfig, read_config
 from .utils.utils import (
     assert_type,
     simple_parse_kwargs_string,
@@ -662,6 +664,63 @@ def load_scores(path: Path) -> Scores:
     return Scores(mmap, info, offsets)
 
 
+def _load_legacy_pt_scores(score_path: str) -> tuple[torch.Tensor, bool]:
+    """Read a bare ``scores.pt`` written by a MAGIC run that predates score
+    directories.
+
+    A 2-D tensor is ambiguous — per-token scores are ``[docs, seq_len]`` and
+    per-query scores are ``[docs, queries]`` — so the run config beside it
+    decides: per-query iff the run used ``query_method: none``. Already in the
+    loss-diff convention, so nothing is negated.
+
+    TODO: Lucia Quirke remove December 2026
+    """
+    scores = torch.load(score_path, map_location="cpu")
+    if not isinstance(scores, torch.Tensor) or scores.ndim not in (2, 3):
+        return scores, False
+    if scores.ndim == 3:
+        return scores, True
+
+    cfg_path = Path(score_path).parent / CONFIG_FILENAME
+    if not cfg_path.is_file():
+        return scores, False
+
+    try:
+        steps = read_config(cfg_path)["steps"]
+    except (ValueError, yaml.YAMLError):
+        return scores, False
+
+    for step in steps:
+        for step_cfg in step.values():
+            if isinstance(step_cfg, dict):
+                return scores, step_cfg.get("query_method") == "none"
+    return scores, False
+
+
+def load_scores_loss_signed(score_path: str) -> tuple[torch.Tensor, bool]:
+    """Loads scores using the sign convention that negative scores reduce
+    query loss (proponents are negative)."""
+    # TODO: Lucia Quirke remove December 2026
+    if score_path.endswith(".pt"):
+        return _load_legacy_pt_scores(score_path)
+
+    loaded = load_scores(Path(score_path))
+    score_cfg = load_subconfig(score_path, "score_cfg", ScoreConfig)
+    negate = score_cfg is not None and score_cfg.higher_is_better
+
+    if loaded.offsets is not None:
+        scores = loaded.to_grid()
+    else:
+        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, loaded.num_scores > 1
+
+
 def sorted_checkpoints(folder: str) -> list[tuple[int, str]]:
     """
     Return a list of (step, filepath) sorted by step

bergson/magic/cli.py

@@ -6,6 +6,7 @@
 from dataclasses import asdict
 from pathlib import Path
 
+import numpy as np
 import torch
 import torch.distributed as dist
 from datasets import Dataset, concatenate_datasets
@@ -28,15 +29,17 @@
 )
 
 from ..config.config import TrainingConfig, ValidationConfig
-from ..config.config_io import read_first_step_config, save_run_config
+from ..config.config_io import save_run_config
+from ..data import compute_num_token_grads, load_scores_loss_signed
 from ..distributed import launch_distributed_run
+from ..score.score_writer import save_sequence_scores, save_token_scores
 from ..utils.load_from_optimizer import (
     save_second_moments_as_optimizer_pt,
 )
 from ..utils.logging import wandb_log_fn
 from ..utils.utils import get_device, get_device_index
 from ..utils.worker_utils import setup_data_pipeline
-from ..validate import load_attribution_scores, validate_scores
+from ..validate import validate_scores
 from .config import MagicConfig
 from .data_stream import DataStream, pad_dataset_to_batch_size
 from .grad_accum import accumulate_grads
@@ -244,20 +247,11 @@ def compute_per_query_magic_scores(
 
 
 def scores_are_per_token(score_path: str) -> bool:
-    if os.path.isdir(score_path):
-        info_path = os.path.join(score_path, "info.json")
-        if not os.path.isfile(info_path):
-            return False
-        with open(info_path) as f:
-            return bool(json.load(f).get("attribute_tokens", 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 == 3 or (scores.ndim == 2 and scores.shape[1] > 1)
-    )
+    info_path = os.path.join(score_path, "info.json")
+    if not os.path.isfile(info_path):
+        return False
+    with open(info_path) as f:
+        return bool(json.load(f).get("attribute_tokens", False))
 
 
 def attach_doc_ids_if_missing(dataset: Dataset) -> Dataset:
@@ -286,16 +280,46 @@ def attach_doc_ids_if_missing(dataset: Dataset) -> Dataset:
     )
 
 
-def save_doc_ids(run_path: str, train_dataset: Dataset, pad_count: int) -> str:
-    """Write ``doc_ids.pt`` beside a per-token ``scores.pt``."""
-    doc_ids = torch.tensor(train_dataset["doc_ids"])
+def save_magic_scores(
+    run_path: str,
+    scores: torch.Tensor,
+    train_dataset: Dataset,
+    pad_count: int,
+    per_token: bool,
+) -> str:
+    """Write MAGIC scores as a score directory under ``<run_path>/scores``.
+
+    Columns past a row's ``length - 1`` never receive gradient, so the dense
+    per-token grid packs into the ragged store without loss.
+    """
+    path = Path(run_path) / "scores"
+
+    if not per_token:
+        arr = scores.float().numpy()
+        save_sequence_scores(path, arr if arr.ndim > 1 else arr[:, None])
+        print(f"Saved attribution scores to {path}")
+        return str(path)
+
+    num_token_grads = compute_num_token_grads(train_dataset)
+    doc_ids = np.asarray(train_dataset["doc_ids"], dtype=np.int64)
     if pad_count:
+        num_token_grads = num_token_grads[:-pad_count]
         doc_ids = doc_ids[:-pad_count]
 
-    doc_ids_path = os.path.join(run_path, "doc_ids.pt")
-    torch.save(doc_ids, doc_ids_path)
-    print(f"Saved doc_ids to {doc_ids_path}")
-    return doc_ids_path
+    offsets = np.zeros(len(num_token_grads) + 1, dtype=np.int64)
+    np.cumsum(num_token_grads, out=offsets[1:])
+
+    grid = scores.float().numpy()
+    if grid.ndim == 2:
+        grid = grid[:, :, None]
+    flat = np.concatenate(
+        [grid[r, : num_token_grads[r]] for r in range(len(num_token_grads))]
+    )
+
+    save_token_scores(path, flat, offsets)
+    np.save(path / "doc_ids.npy", doc_ids)
+    print(f"Saved attribution scores to {path}")
+    return str(path)
 
 
 def shuffled_epochs(dataset: Dataset, seed: int, num_epochs: int) -> Dataset:
@@ -512,9 +536,9 @@ def worker(
         if global_rank == 0:
             print(f"Baseline loss: {baseline}")
             print(f"Score summary: {describe(scores.flatten())}")
-            score_path = os.path.join(run_cfg.run_path, "scores.pt")
-            torch.save(scores, score_path)
-            print(f"Saved per-query attribution scores to {score_path}")
+            score_path = save_magic_scores(
+                run_cfg.run_path, scores, train_dataset, pad_count, bool(per_token)
+            )
     elif not score_path:
         # Sanity check
         if not isinstance(run_cfg, MagicConfig):
@@ -566,16 +590,11 @@ def worker(
             summ = describe(scores.flatten())
             print(f"Score summary: {summ}")
 
-            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):
-        scores, multi_query = load_attribution_scores(score_path)
+            score_path = save_magic_scores(
+                run_cfg.run_path, scores, train_dataset, pad_count, bool(per_token)
+            )
     else:
-        scores = torch.load(score_path, map_location="cpu")
-
-    if per_token and global_rank == 0:
-        save_doc_ids(run_cfg.run_path, train_dataset, pad_count)
+        scores, multi_query = load_scores_loss_signed(score_path)
 
     stream.requires_grad = False
 

bergson/validate.py

@@ -29,60 +29,16 @@
     set_verbosity_error as hf_set_verbosity_error,
 )
 
-from .config.config import ScoreConfig, ValidationConfig
-from .config.config_io import load_subconfig, read_first_step_config, save_run_config
-from .data import Scores, load_scores, pad_and_tensor
+from .config.config import ValidationConfig
+from .config.config_io import save_run_config
+from .data import load_scores_loss_signed, pad_and_tensor
 from .magic.data_stream import DataStream, pad_dataset_to_batch_size
 from .magic.trainer import TrainerState, prepare_trainer
 from .utils.csv_writer import CSVWriter
 from .utils.utils import get_device, simple_parse_kwargs_string
 from .utils.worker_utils import setup_data_pipeline
 
 
-def load_attribution_scores(score_path: str) -> tuple[torch.Tensor, bool]:
-    """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 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.
-    """
-    if os.path.isdir(score_path):
-        loaded = load_scores(Path(score_path))
-        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, 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, loaded.num_scores > 1
-
-    scores = torch.load(score_path, map_location="cpu")
-    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(
     run_cfg: ValidationConfig, multi_query: bool, num_subsets: int
 ) -> str:
@@ -572,7 +528,7 @@ def evaluate_retrained(
     subsets = [torch.tensor(s, dtype=torch.long) for s in subset_lists]
 
     # Load per-query attribution scores (mirrors run_magic's score loading).
-    scores, multi_query = load_attribution_scores(score_path)
+    scores, multi_query = load_scores_loss_signed(score_path)
     if not multi_query:
         assert (
             scores.ndim == 1 or scores.shape[1] == 1

docs/magic.rst

@@ -32,25 +32,35 @@ Output files
 
 After a run completes, ``run_cfg.run_path`` contains:
 
-* ``scores.pt`` — attribution scores tensor. Shape depends on the weight
-  parameterization:
-
-  * Per-example (1D weights): ``(num_train_docs,)``, indexed directly by
-    ``doc_id``.
-  * Per-token (2D weights): ``(num_chunks, seq_len)``, indexed by
-    ``(chunk_idx, token_idx)`` in the *post-shuffle* order used during
-    training. Pad rows appended to make the dataset divisible by
-    ``batch_size`` are trimmed before saving.
-
-* ``doc_ids.pt`` — written alongside ``scores.pt`` for every per-token
-  run, shape ``(num_chunks, seq_len)`` matching ``scores.pt`` row-for-row.
-  Each entry is the original (pre-shuffle) document id for that token
-  position. Downstream aggregation is one line:
+* ``scores/`` — a score directory, the same self-describing format the
+  scoring pipeline writes. ``info.json`` records ``attribute_tokens`` and
+  ``num_scores``, so consumers never infer the layout from the shape.
+  Read it with :func:`bergson.data.load_scores_loss_signed`, which
+  returns ``(scores, multi_query)``:
+
+  * Per-example: ``(num_train_docs, 1)``, indexed directly by ``doc_id``.
+  * Per-token: ``(num_chunks, seq_len)``, indexed by ``(chunk_idx,
+    token_idx)`` in the *post-shuffle* order used during training.
+  * Per-query (``query_method: none``) adds a trailing query axis, so
+    per-token per-query scores are ``(num_chunks, seq_len,
+    num_query_docs)``.
+
+  Pad rows appended to make the dataset divisible by ``batch_size`` are
+  trimmed before saving. Per-token scores are stored ragged — a row holds
+  ``length - 1`` values, the positions ``weighted_causal_lm_ce`` can reach
+  — and are unpacked back into the dense grid on load.
+
+* ``scores/doc_ids.npy`` — written for every per-token run, shape
+  ``(num_chunks, seq_len)`` matching the loaded scores row-for-row. Each
+  entry is the original (pre-shuffle) document id for that token position.
+  Downstream aggregation is one line:
 
   .. code-block:: python
 
-     scores = torch.load("scores.pt")      # (num_chunks, seq_len)
-     doc_ids = torch.load("doc_ids.pt")    # (num_chunks, seq_len)
+     from bergson.data import load_scores_loss_signed
+
+     scores, _ = load_scores_loss_signed("runs/magic/scores")
+     doc_ids = torch.from_numpy(np.load("runs/magic/scores/doc_ids.npy"))
      num_docs = int(doc_ids.max()) + 1
      per_doc = torch.zeros(num_docs, dtype=scores.dtype)
      per_doc.scatter_add_(0, doc_ids.flatten(), scores.flatten())