Blame for code/scripts/dedup_arm_b_shards.py

"""One-shot, surgical dedup of the Arm B preference shards.

WHY THIS EXISTS. The ref-logprob cache build refused the Arm B shards with
`ValueError: duplicate pair_id` -- correctly. The train split held 11,404
rows with 87 pair_ids present exactly twice, and held_out held 596 rows with
2 doubled ids; every second occurrence is byte-identical to its first
(same input_ids, boundaries, rejected side, and extras). Root cause: PRM800K
carries multiple annotation records per problem, so two records can yield the
IDENTICAL pair site; `PairCapper.offer` caps per (problem_id, step_index) but
never checks pair_id, and `curate()` merges everything per problem
(nmo/data/prm800k_pairs.py, `by_problem.setdefault(...).append(p)`) without a
seen-pair_id check -- so the depth-round top-up re-selects a pair the first
round already took.

WHAT IT DOES. For each split (train, held_out):
  1. read the live shards via the store's own reader;
  2. keep the FIRST occurrence of every pair_id, drop second-and-later
     occurrences -- REFUSING unless the dropped row is content-identical to
     the kept one (a non-identical duplicate is a different bug and must not
     be silently resolved by row order);
  3. write fresh shards through the same ShardWriter/receipt API the build
     used, preserving shard naming (`arm_b-train-NNNN` etc.).
Then it asserts zero train/held pair_id overlap, atomically swaps
`arm_b -> arm_b_predup` and `arm_b_dedup -> arm_b` (the predup directory is
KEPT as evidence), and updates data/manifest.json in place:
`counts.arm_b_curated`, the two `arm_b-*` shard hashes (same
content_hash-of-receipt-hashes chain `_ShardSet.finalize` uses), and the
`length_bias_baseline` Arm B entries re-measured over the deduped TRAIN split
with the exact function scripts/dpo.py uses
(align.train.dpo_stage.measure_chosen_longer_fraction). A provenance block
`arm_b_dedup` records old -> new for everything it touched.

Deterministic (pure function of the input shards; no RNG), CPU-only.
The ref-logprob caches keyed on the old manifest hashes are stale after this
and must be re-precomputed -- that is the point.

Run:

    cd /apps/development/rlhf/nano-mistral-o1
    CUDA_VISIBLE_DEVICES= .venv/bin/python scripts/dedup_arm_b_shards.py --dry-run
    CUDA_VISIBLE_DEVICES= .venv/bin/python scripts/dedup_arm_b_shards.py

Refuses to run twice: an existing `arm_b_predup` means the swap already
happened.
"""

from __future__ import annotations

import os

# CPU-only by construction: never let a lazily-imported torch see a GPU.
os.environ.setdefault("CUDA_VISIBLE_DEVICES", "")

import argparse
import json
import logging
from pathlib import Path

import numpy as np

logger = logging.getLogger("dedup_arm_b")

REPO = Path("/apps/development/rlhf/nano-mistral-o1")
SPLITS = (("train", "arm_b-train"), ("held_out", "arm_b-held"))
#: manifest hash key per split prefix, matching build_all's `hashes` dict.
MANIFEST_HASH_KEY = {"arm_b-train": "arm_b-train", "arm_b-held": "arm_b-held"}


def rows_identical(a, b) -> bool:
    """Full content identity between two store Rows -- every stored field.

    Anything short of this (say, comparing input_ids only) would let a real
    divergence hide behind a shared pair_id and get resolved by row order.
    """
    return (
        np.array_equal(a.input_ids, b.input_ids)
        and np.array_equal(a.boundary_positions, b.boundary_positions)
        and a.problem_id == b.problem_id
        and a.subject == b.subject
        and sorted(a.ragged) == sorted(b.ragged)
        and all(np.array_equal(a.ragged[f], b.ragged[f]) for f in a.ragged)
        and a.extras == b.extras
    )


def dedup_row_indices(ds) -> tuple[list[int], list[int]]:
    """(kept, dropped) row indices: first occurrence per pair_id wins.

    Raises if any dropped row is NOT content-identical to its kept first
    occurrence -- that would be a semantic collision, not a re-added pair,
    and this script is only licensed to remove exact re-adds.
    """
    first_by_pair: dict[str, int] = {}
    kept: list[int] = []
    dropped: list[int] = []
    for i in range(len(ds)):
        row = ds[i]
        pair_id = row.extras["pair_id"]
        if pair_id not in first_by_pair:
            first_by_pair[pair_id] = i
            kept.append(i)
            continue
        j = first_by_pair[pair_id]
        if not rows_identical(ds[j], row):
            raise ValueError(
                f"pair_id {pair_id!r} occurs at rows {j} and {i} with "
                "DIFFERENT content. That is a hash/semantics collision, not a "
                "duplicate re-add; refusing to pick a winner by row order."
            )
        dropped.append(i)
    return kept, dropped


def write_split(ds, kept: list[int], out_dir: Path, prefix: str) -> str:
    """Write the kept rows through the build's own shard-set API.

    Returns the manifest-level content hash (content_hash over the per-shard
    receipt hashes -- the same chain `_ShardSet.finalize` produces, so the
    updated manifest entry is computed by the code that defined it).
    """
    from align.data.build import _PREF_FIELDS, _ShardSet

    shard_set = _ShardSet(out_dir, prefix, _PREF_FIELDS)
    for i in kept:
        row = ds[i]
        shard_set.add(
            input_ids=row.input_ids,
            problem_id=row.problem_id,
            subject=row.subject,
            boundary_positions=row.boundary_positions,
            ragged={
                "rejected_input_ids": row.ragged["rejected_input_ids"],
                "rejected_boundary_positions": row.ragged["rejected_boundary_positions"],
            },
            **row.extras,
        )
    return shard_set.finalize()


def step_basis_fraction(rows) -> float:
    """Arm B's step-basis chosen-longer fraction over collator-shaped rows.

    Mirrors build.py's `chosen_longer_fraction_step_tokens`: length measured
    from the SHARED PREFIX (the one step that differs), counted over every
    row. Kept next to the completion-basis measurement so the manifest's two
    Arm B baselines stay siblings of the same dataset state.
    """
    longer = sum(
        int(
            len(r["chosen_input_ids"]) - int(r["shared_prefix_len"])
            > len(r["rejected_input_ids"]) - int(r["shared_prefix_len"])
        )
        for r in rows
    )
    return longer / len(rows) if rows else float("nan")


def main(argv=None) -> int:
    logging.basicConfig(level=logging.INFO, format="%(levelname)s %(name)s: %(message)s")
    ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    ap.add_argument("--root", type=Path, default=REPO, help="repo root")
    ap.add_argument(
        "--dry-run", action="store_true", help="measure and report; write and swap nothing"
    )
    args = ap.parse_args(argv)

    from align.data.loader import preference_dataset
    from align.data.store import ShardDataset

    shards_root = args.root / "data" / "shards"
    live = shards_root / "arm_b"
    staged = shards_root / "arm_b_dedup"
    predup = shards_root / "arm_b_predup"
    manifest_path = args.root / "data" / "manifest.json"

    if predup.exists():
        raise SystemExit(
            f"{predup} already exists: the swap already happened. This script "
            "is one-shot; a second run would dedup the already-deduped shards "
            "and overwrite the evidence."
        )
    if staged.exists():
        raise SystemExit(
            f"{staged} already exists (a previous run died mid-write?). "
            "Inspect and delete it before re-running; a half-written shard "
            "set masquerades as resumable."
        )

    # -- pass 1: dedup into the staging directory ---------------------------
    pair_ids: dict[str, set[str]] = {}
    before_after: dict[str, tuple[int, int]] = {}
    new_hashes: dict[str, str] = {}
    for split, prefix in SPLITS:
        ds = ShardDataset(live / split)
        kept, dropped = dedup_row_indices(ds)
        before_after[split] = (len(ds), len(kept))
        pair_ids[split] = {ds[i].extras["pair_id"] for i in kept}
        logger.info(
            "%s: %d rows -> %d kept (%d exact-duplicate re-adds dropped)",
            split,
            len(ds),
            len(kept),
            len(dropped),
        )
        if not args.dry_run:
            new_hashes[prefix] = write_split(ds, kept, staged / split, prefix)

    overlap = pair_ids["train"] & pair_ids["held_out"]
    if overlap:
        raise SystemExit(
            f"{len(overlap)} pair_ids appear in BOTH deduped splits (e.g. "
            f"{sorted(overlap)[:3]}); the train/held split is broken and a "
            "dedup must not paper over it."
        )

    if args.dry_run:
        logger.info("dry run: nothing written, nothing swapped")
        return 0

    # -- pass 2: verify the staged shards through a fresh reader ------------
    for split, _prefix in SPLITS:
        staged_ds = ShardDataset(staged / split)
        staged_ids = [staged_ds[i].extras["pair_id"] for i in range(len(staged_ds))]
        if len(staged_ids) != len(set(staged_ids)) or len(staged_ids) != before_after[split][1]:
            raise SystemExit(
                f"staged {split} re-read: {len(staged_ids)} rows, "
                f"{len(set(staged_ids))} unique -- the write did not produce "
                "what the dedup computed; refusing to swap."
            )

    # -- pass 3: atomic swap; the pre-dedup shards are KEPT as evidence -----
    live.rename(predup)
    staged.rename(live)
    logger.info("swapped: %s -> %s, %s -> %s", "arm_b", "arm_b_predup", "arm_b_dedup", "arm_b")

    # -- pass 4: manifest update, through the numbers' own definitions ------
    manifest = json.loads(manifest_path.read_text())
    old = {
        "counts.arm_b_curated": manifest["counts"]["arm_b_curated"],
        "hashes.arm_b-train": manifest["hashes"]["arm_b-train"],
        "hashes.arm_b-held": manifest["hashes"]["arm_b-held"],
        "arm_b_chosen_longer_fraction": manifest["length_bias_baseline"][
            "arm_b_chosen_longer_fraction"
        ],
        "arm_b_chosen_longer_fraction_step_basis": manifest["length_bias_baseline"][
            "arm_b_chosen_longer_fraction_step_basis"
        ],
    }

    # Re-measure the length-bias baseline over the deduped TRAIN split with
    # the exact function scripts/dpo.py falls back to, on the exact rows it
    # would build (read-only imports; nothing here re-implements the metric).
    from align.train.dpo_stage import measure_chosen_longer_fraction

    rows = preference_dataset(live / "train")
    new_fraction = measure_chosen_longer_fraction(rows)
    new_fraction_step = step_basis_fraction(rows)

    n_train, n_held = before_after["train"][1], before_after["held_out"][1]
    manifest["counts"]["arm_b_curated"] = n_train + n_held
    manifest["hashes"]["arm_b-train"] = new_hashes["arm_b-train"]
    manifest["hashes"]["arm_b-held"] = new_hashes["arm_b-held"]
    manifest["length_bias_baseline"]["arm_b_chosen_longer_fraction"] = new_fraction
    manifest["length_bias_baseline"]["arm_b_chosen_longer_fraction_step_basis"] = new_fraction_step
    manifest["arm_b_dedup"] = {
        "reason": (
            "87 train + 2 held_out pair_ids were present exactly twice, every "
            "second occurrence content-identical to its first (curation "
            "depth-round top-up re-added pairs already selected; the merge in "
            "prm800k_pairs.curate never deduped by pair_id). The ref-logprob "
            "cache build correctly refused the duplicates."
        ),
        "script": "scripts/dedup_arm_b_shards.py",
        "train_rows": {"before": before_after["train"][0], "after": n_train},
        "held_out_rows": {"before": before_after["held_out"][0], "after": n_held},
        "old": old,
        "predup_shards": str(predup),
    }
    manifest_path.write_text(json.dumps(manifest, indent=1))

    logger.info("train: %d -> %d rows", before_after["train"][0], n_train)
    logger.info("held_out: %d -> %d rows", before_after["held_out"][0], n_held)
    logger.info("counts.arm_b_curated: %s -> %d", old["counts.arm_b_curated"], n_train + n_held)
    logger.info(
        "arm_b_chosen_longer_fraction: %s -> %.6f (step basis %s -> %.6f)",
        old["arm_b_chosen_longer_fraction"],
        new_fraction,
        old["arm_b_chosen_longer_fraction_step_basis"],
        new_fraction_step,
    )
    logger.info(
        "hashes: arm_b-train %s -> %s, arm_b-held %s -> %s",
        old["hashes.arm_b-train"],
        new_hashes["arm_b-train"],
        old["hashes.arm_b-held"],
        new_hashes["arm_b-held"],
    )
    logger.info(
        "NOTE: ref caches keyed on the old manifest hashes are now stale; "
        "re-run scripts/precompute_ref_logprobs.py before scripts/dpo.py."
    )
    return 0


if __name__ == "__main__":
    # House rule: everything above is import-safe because spawn-based worker
    # pools re-import __main__; the entry point stays behind this guard.
    raise SystemExit(main())