From 3fabb2a98f7ccd074f876c14660f60b32b061b2b Mon Sep 17 00:00:00 2001 From: drewOrc <36374426+drewOrc@users.noreply.github.com> Date: Wed, 23 Sep 2026 14:44:16 +0800 Subject: [PATCH 1/4] Add k-shot sampling, step planning, pilots, curves and baselines Step 3 code, no training runs. The learning-curve protocol lives in configs/curve.yaml and is filled in by hand from the pilots. - sampling: k rows per intent, oos rows from a hardcoded ceil(2.5k) table; k=100 returns the whole train split in its original order. - steps: train for max(min_train_steps, epoch steps); train() checks the steps it ran against the plan and records which rule decided. - runs: resume rules shared by AC2, pilots and curves, plus equivalent_run, which lets an existing run stand in for a curve point only when every field, the rows, the step count and the seed agree. The committed AC2 records are equivalent to BERT k=100 at lr 5e-5. - pilots: lr and S_min pilots on validation only; the scoring and prediction functions raise LeakageError on any other split. - curves: 6 k x 3 seeds per encoder and the ModernBERT OOS ablation, reusing AC2 at k=100 when equivalent; no curve weights are kept. - baselines: majority class and TF-IDF centroid on the same sample as the encoders, archived in the encoder format. - archive format 2 records the SHA-256 of the three split files; format 1 archives are back-filled on read only at the pinned revision. - metrics: Wilson score interval. --- Makefile | 29 ++- configs/curve.yaml | 15 ++ configs/modernbert-base.yaml | 18 ++ src/tinyrouter/ac2.py | 116 +++--------- src/tinyrouter/archive.py | 128 +++++++++++++- src/tinyrouter/baselines.py | 267 ++++++++++++++++++++++++++++ src/tinyrouter/config.py | 48 ++++- src/tinyrouter/curves.py | 149 ++++++++++++++++ src/tinyrouter/evaluate.py | 17 +- src/tinyrouter/metrics.py | 21 +++ src/tinyrouter/pilots.py | 329 +++++++++++++++++++++++++++++++++++ src/tinyrouter/protocol.py | 141 +++++++++++++++ src/tinyrouter/runs.py | 225 ++++++++++++++++++++++++ src/tinyrouter/sampling.py | 97 +++++++++++ src/tinyrouter/steps.py | 57 ++++++ src/tinyrouter/train.py | 48 ++++- tests/archive_fakes.py | 5 +- tests/run_fakes.py | 112 ++++++++++++ tests/test_archive.py | 106 +++++++++++ tests/test_baselines.py | 129 ++++++++++++++ tests/test_config.py | 28 ++- tests/test_curves.py | 194 +++++++++++++++++++++ tests/test_metrics.py | 40 +++++ tests/test_pilots.py | 218 +++++++++++++++++++++++ tests/test_protocol.py | 113 ++++++++++++ tests/test_runs.py | 142 +++++++++++++++ tests/test_sampling.py | 126 ++++++++++++++ tests/test_steps.py | 64 +++++++ tests/test_train_smoke.py | 37 ++++ 29 files changed, 2904 insertions(+), 115 deletions(-) create mode 100644 configs/curve.yaml create mode 100644 configs/modernbert-base.yaml create mode 100644 src/tinyrouter/baselines.py create mode 100644 src/tinyrouter/curves.py create mode 100644 src/tinyrouter/pilots.py create mode 100644 src/tinyrouter/protocol.py create mode 100644 src/tinyrouter/runs.py create mode 100644 src/tinyrouter/sampling.py create mode 100644 src/tinyrouter/steps.py create mode 100644 tests/run_fakes.py create mode 100644 tests/test_baselines.py create mode 100644 tests/test_curves.py create mode 100644 tests/test_pilots.py create mode 100644 tests/test_protocol.py create mode 100644 tests/test_runs.py create mode 100644 tests/test_sampling.py create mode 100644 tests/test_steps.py diff --git a/Makefile b/Makefile index f250813..3a36775 100644 --- a/Makefile +++ b/Makefile @@ -1,4 +1,5 @@ -.PHONY: setup lint format test test-network smoke train evaluate ac2 verify-logits report clean-checkpoints +.PHONY: setup lint format test test-network smoke train evaluate ac2 pilot-lr pilot-steps baselines \ + curve oos-ablation verify-logits report clean-checkpoints CONFIG ?= configs/bert-base.yaml SEED ?= 42 @@ -64,6 +65,32 @@ evaluate: ac2: uv run $(UV_ENV) python -m tinyrouter.ac2 --config configs/bert-base.yaml $(if $(filter 1,$(FORCE)),--force,) +# Step 3 (docs/PLAN.md section 4, hyperparameter protocol). Order: +# make pilot-lr then copy each model's selected lr into configs/curve.yaml +# make pilot-steps then copy the selected S_min into configs/curve.yaml +# make curve MODEL=bert ; make curve MODEL=modernbert ; make oos-ablation +# Pilots read and write validation numbers only (results/pilots/*.json) and +# never edit configs/curve.yaml. Curves and the ablation resume like ac2 and +# keep no weights, only logits. `curve` runs the cheap baselines first. +# `make baselines` runs them alone. +pilot-lr: + uv run $(UV_ENV) python -m tinyrouter.pilots lr + +pilot-steps: + uv run $(UV_ENV) python -m tinyrouter.pilots steps + +# Majority-class and TF-IDF centroid baselines on every (k, seed) sample; seconds. +baselines: + uv run python -m tinyrouter.baselines + +# Checks configs/curve.yaml is filled in, then runs the baselines, then the curve. +curve: + @case "$(MODEL)" in bert|modernbert) ;; *) echo "usage: make curve MODEL=bert|modernbert"; exit 2;; esac + uv run $(UV_ENV) python -m tinyrouter.curves --model $(MODEL) + +oos-ablation: + uv run $(UV_ENV) python -m tinyrouter.curves --ablation + # Every archive listed in results/logits-manifest.json is present and matches its SHA-256. verify-logits: uv run python -m tinyrouter.archive diff --git a/configs/curve.yaml b/configs/curve.yaml new file mode 100644 index 0000000..e9a7aaf --- /dev/null +++ b/configs/curve.yaml @@ -0,0 +1,15 @@ +# Learning-curve hyperparameters (docs/PLAN.md section 4, hyperparameter +# protocol). Filled in by hand from the pilot outputs: +# make pilot-lr -> results/pilots/lr.json -> models..learning_rate +# make pilot-steps -> results/pilots/steps.json -> min_train_steps +# null means that pilot has not been run; `make curve` refuses to start. +# Allowed values: learning_rate in {1.0e-5, 2.0e-5, 5.0e-5} (write the +# ".0": YAML reads 5e-5 as a string), min_train_steps in {100, 200, 400}. +min_train_steps: null +models: + bert: + config: configs/bert-base.yaml + learning_rate: null + modernbert: + config: configs/modernbert-base.yaml + learning_rate: null diff --git a/configs/modernbert-base.yaml b/configs/modernbert-base.yaml new file mode 100644 index 0000000..395e106 --- /dev/null +++ b/configs/modernbert-base.yaml @@ -0,0 +1,18 @@ +# answerdotai/ModernBERT-base, the main encoder (docs/PLAN.md section 4). +# Revision checked against the Hub API on 2026-09-23 (Apache-2.0). +# Everything except the model matches bert-base.yaml, so the two curves +# differ in the encoder only. +model_name: answerdotai/ModernBERT-base +model_revision: 8949b909ec900327062f0ebf497f51aef5e6f0c8 +seed: 42 +per_intent: null +max_length: 64 +# Placeholder: chosen by `make pilot-lr` from {1e-5, 2e-5, 5e-5} on +# validation only, then written into configs/curve.yaml. Training refuses null. +learning_rate: null +weight_decay: 0.01 +warmup_ratio: 0.1 +num_train_epochs: 5 +train_batch_size: 32 +eval_batch_size: 128 +device: auto diff --git a/src/tinyrouter/ac2.py b/src/tinyrouter/ac2.py index dbfbc23..99eb66c 100644 --- a/src/tinyrouter/ac2.py +++ b/src/tinyrouter/ac2.py @@ -10,11 +10,12 @@ printout list them for that purpose. The test numbers are the verdict and are not a tuning signal. -Resumable. A seed is skipped only when its results JSON was produced by -the current config (output paths aside) and its logits archive has the -same SHA-256 in the results JSON, in the manifest and on disk. Trained -weights are reused only when their ``train_summary.json`` records the -current config. Anything else is cleared and redone. ``--force`` clears +Resumable, with the rules in runs.py: a seed is skipped only when its +results JSON was produced by the current config (output paths aside) and +its logits archive has the same SHA-256 in the results JSON, in the +manifest and on disk. Trained weights are reused only when their +``train_summary.json`` records the current config. Anything else is +cleared and redone. ``--force`` clears every seed's outputs before training anything, so a crash halfway never leaves old and new runs mixed. Weights of every seed except 42 are deleted once that seed's logits are archived (disk is tight; seed 42 is kept as @@ -25,14 +26,22 @@ import argparse import json -import shutil -from collections.abc import Callable from pathlib import Path -from tinyrouter.archive import ArchiveError, load_logits, read_manifest, remove_from_manifest from tinyrouter.config import RunConfig, load_config -from tinyrouter.data import sha256_of from tinyrouter.evaluate import RunPaths +from tinyrouter.runs import ( + EvaluateFn, + Log, + TrainFn, + clear_outputs, + default_evaluate, + default_train, + is_done, + reusable_weights, + run_dir, + run_one, +) SEEDS = (42, 43, 44) KEEP_WEIGHTS_SEEDS = frozenset({42}) @@ -42,10 +51,16 @@ EXPECTED_OOS_TRAIN_ROWS = 250 EXPECTED_TEST_ROWS = 5_500 REPORTED = ("in_scope_accuracy_150", "oos_recall_151", "accuracy_8", "oos_recall_8") - -TrainFn = Callable[[RunConfig], Path] -EvaluateFn = Callable[[RunConfig, Path], dict[str, object]] -Log = Callable[[str], None] +# Re-exported: the resume rules live in runs.py and are shared with the curves. +__all__ = [ + "clear_outputs", + "default_evaluate", + "default_train", + "is_done", + "reusable_weights", + "run_ac2", + "run_dir", +] class SetupError(ValueError): @@ -56,69 +71,8 @@ def expected_run_name(seed: int) -> str: return f"bert-base-uncased-full-seed{seed}" -def run_dir(config: RunConfig) -> Path: - return Path(config.checkpoint_root) / config.run_name - - -def is_done(config: RunConfig) -> bool: - """This config's results JSON exists and its archive is intact and the one it was scored on.""" - paths = RunPaths.of(config) - if not (paths.results_json.exists() and paths.logits.exists()): - return False - record = json.loads(paths.results_json.read_text(encoding="utf-8")) - if not config.matches(record.get("config")): - return False - recorded = record.get("logits", {}).get("sha256") - listed = read_manifest(paths.manifest).get(paths.logits.name, {}).get("sha256") - if not recorded or not recorded == listed == sha256_of(paths.logits): - return False - try: - archive = load_logits(paths.logits) - except ArchiveError: - return False - return archive.metadata["seed"] == config.seed - - -def reusable_weights(config: RunConfig) -> bool: - """Final weights exist and were trained with exactly this config.""" - final = run_dir(config) / "final" - summary = final / "train_summary.json" - if not ((final / "model.safetensors").exists() and summary.exists()): - return False - return config.matches(json.loads(summary.read_text(encoding="utf-8")).get("config")) - - -def clear_outputs(config: RunConfig, log: Log) -> None: - """Remove this seed's weights, results JSON, logits archive and manifest entry.""" - paths = RunPaths.of(config) - removed = [str(p) for p in (paths.results_json, paths.logits) if p.exists()] - for name in removed: - Path(name).unlink() - if remove_from_manifest(paths.manifest, paths.logits.name): - removed.append(f"manifest entry {paths.logits.name}") - if run_dir(config).exists(): - shutil.rmtree(run_dir(config)) - removed.append(str(run_dir(config))) - if removed: - log(f"seed {config.seed}: cleared {', '.join(removed)}") - - def run_seed(config: RunConfig, train_fn: TrainFn, evaluate_fn: EvaluateFn, log: Log) -> None: - if is_done(config): - log(f"seed {config.seed}: results and logits for this config already archived, skipping") - else: - if reusable_weights(config): - log(f"seed {config.seed}: reusing weights trained with this config") - final_dir = run_dir(config) / "final" - else: - clear_outputs(config, log) - log(f"seed {config.seed}: training {config.run_name}") - final_dir = train_fn(config) - evaluate_fn(config, final_dir) - log(f"seed {config.seed}: evaluated, logits archived") - if config.seed not in KEEP_WEIGHTS_SEEDS and run_dir(config).exists(): - shutil.rmtree(run_dir(config)) - log(f"seed {config.seed}: deleted weights {run_dir(config)} (only seed 42 is kept)") + run_one(config, train_fn, evaluate_fn, log, keep_weights=config.seed in KEEP_WEIGHTS_SEEDS) def check_setup(record: dict[str, object], seed: int) -> None: @@ -209,18 +163,6 @@ def run_ac2( return result -def default_train(config: RunConfig) -> Path: - from tinyrouter.train import prepare_train_split, train - - return train(config, prepare_train_split(config), run_dir(config)) - - -def default_evaluate(config: RunConfig, model_dir: Path) -> dict[str, object]: - from tinyrouter.evaluate import evaluate - - return evaluate(config, model_dir) - - def main(argv: list[str] | None = None) -> None: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--config", default="configs/bert-base.yaml") diff --git a/src/tinyrouter/archive.py b/src/tinyrouter/archive.py index 7749572..2d8c6ba 100644 --- a/src/tinyrouter/archive.py +++ b/src/tinyrouter/archive.py @@ -9,6 +9,20 @@ is recorded in ``results/logits-manifest.json``, which is committed, so a copy downloaded from a GitHub Release can be checked byte for byte against what the run produced. + +Format versions (``format_version`` in the metadata): + +- 1 (AC2 archives, 2026-09-23): no per-file dataset checksums, only + ``dataset_revision``. +- 2: adds ``dataset_sha256``, the SHA-256 of the train, validation and + test parquet files the run read (each verified against ``SPLIT_FILES`` + before use). Only format 2 is written. + +``load_logits`` still reads format 1. It fills ``dataset_sha256`` from the +pinned ``SPLIT_FILES`` when the archive's ``dataset_revision`` is the +revision this code pins (the files at a revision are immutable, and +``download_split_file`` refuses any other bytes), marks the value as +back-filled, and refuses a format-1 archive from any other revision. """ from __future__ import annotations @@ -24,10 +38,15 @@ import numpy as np from tinyrouter.calibrate import SplitLogits -from tinyrouter.data import SplitName, sha256_of +from tinyrouter.data import DATASET_REVISION, SPLIT_FILES, SplitName, sha256_of from tinyrouter.labels import load_label_space -FORMAT_VERSION = 1 +FORMAT_VERSION = 2 +READABLE_FORMAT_VERSIONS = (1, 2) +# The manifest's own layout has not changed since it was introduced. +MANIFEST_FORMAT_VERSION = 1 +DATASET_SPLITS: tuple[SplitName, ...] = ("train", "validation", "test") +BACKFILL_NOTE = "back-filled from SPLIT_FILES when read (format 1 archive)" SPLITS: tuple[SplitName, ...] = ("validation", "test") LOGITS_DTYPE = np.float32 MANIFEST_NAME = "logits-manifest.json" @@ -50,6 +69,14 @@ "created_at": str, } NULLABLE = {"per_intent", "eval_per_intent"} +# Added in format 2: file checksums, and the k-shot size (null unless a curve run). +DATASET_SHA_KEY = "dataset_sha256" +K_SHOT_KEY = "k_shot" + + +def pinned_dataset_sha256() -> dict[str, str]: + """SHA-256 of each split's parquet file at the pinned revision.""" + return {name: SPLIT_FILES[name].sha256 for name in DATASET_SPLITS} class ArchiveError(ValueError): @@ -100,6 +127,64 @@ def utc_now() -> str: def validate_metadata(metadata: dict[str, object]) -> None: + """Refuse metadata that is not a complete, current (format 2) record.""" + validate_common_metadata(metadata) + if metadata["format_version"] != FORMAT_VERSION: + raise ArchiveError( + f"format_version {metadata['format_version']} cannot be written (current is " + f"{FORMAT_VERSION})" + ) + validate_format2_fields(metadata) + + +def validate_format2_fields(metadata: dict[str, object]) -> None: + if K_SHOT_KEY not in metadata: + raise ArchiveError(f"metadata is missing [{K_SHOT_KEY!r}]") + k_shot = metadata[K_SHOT_KEY] + if k_shot is not None and (isinstance(k_shot, bool) or not isinstance(k_shot, int)): + raise ArchiveError(f"metadata[{K_SHOT_KEY!r}] = {k_shot!r} is not an int or null") + validate_dataset_sha256(metadata) + + +def validate_dataset_sha256(metadata: dict[str, object]) -> None: + recorded = metadata.get(DATASET_SHA_KEY) + if not isinstance(recorded, dict) or set(recorded) != set(DATASET_SPLITS): + raise ArchiveError( + f"metadata[{DATASET_SHA_KEY!r}] must map exactly {DATASET_SPLITS} to SHA-256 digests" + ) + if metadata["dataset_revision"] != DATASET_REVISION or recorded != pinned_dataset_sha256(): + raise ArchiveError( + f"dataset (revision {metadata['dataset_revision']}, files {recorded}) is not the " + f"pinned CLINC150 revision {DATASET_REVISION}; the analysis code assumes those files" + ) + + +def upgrade_metadata(metadata: dict[str, object]) -> dict[str, object]: + """Metadata of a readable archive in the current format; format 1 gets dataset_sha256.""" + validate_common_metadata(metadata) + version = metadata["format_version"] + if version == FORMAT_VERSION: + validate_format2_fields(metadata) + return metadata + for key in (DATASET_SHA_KEY, K_SHOT_KEY): + if key in metadata: + raise ArchiveError(f"format {version} archive already carries {key!r}") + if metadata["dataset_revision"] != DATASET_REVISION: + raise ArchiveError( + f"format {version} archive from dataset revision {metadata['dataset_revision']}; " + f"its file checksums can only be back-filled for the pinned revision " + f"{DATASET_REVISION}" + ) + return { + **metadata, + # Format 1 predates k-shot sampling: those runs used per_intent only. + K_SHOT_KEY: None, + DATASET_SHA_KEY: pinned_dataset_sha256(), + "dataset_sha256_source": BACKFILL_NOTE, + } + + +def validate_common_metadata(metadata: dict[str, object]) -> None: missing = sorted(set(REQUIRED_METADATA) - set(metadata)) if missing: raise ArchiveError(f"metadata is missing {missing}") @@ -113,10 +198,10 @@ def validate_metadata(metadata: dict[str, object]) -> None: raise ArchiveError(f"metadata[{key!r}] = {value!r} is not {expected}") if isinstance(value, str) and not value.strip(): raise ArchiveError(f"metadata[{key!r}] is empty") - if metadata["format_version"] != FORMAT_VERSION: + if metadata["format_version"] not in READABLE_FORMAT_VERSIONS: raise ArchiveError( - f"format_version {metadata['format_version']} is not supported (expected " - f"{FORMAT_VERSION})" + f"format_version {metadata['format_version']} is not supported (readable: " + f"{READABLE_FORMAT_VERSIONS})" ) expected_hash = load_label_space().sha256 if metadata["label_space_sha256"] != expected_hash: @@ -167,7 +252,10 @@ def save_logits(path: Path, splits: dict[str, SplitLogits], metadata: dict[str, def load_logits(path: Path) -> LogitsArchive: - """Read an archive back, refusing it if metadata or shapes are incomplete or inconsistent.""" + """Read an archive back, refusing it if metadata or shapes are incomplete or inconsistent. + + The returned metadata is in the current format (see ``upgrade_metadata``). + """ expected_keys = {"metadata"} | {f"{s}_{part}" for s in SPLITS for part in ("logits", "labels")} with np.load(path, allow_pickle=False) as npz: missing = sorted(expected_keys - set(npz.files)) @@ -177,7 +265,7 @@ def load_logits(path: Path) -> LogitsArchive: arrays = {key: npz[key] for key in expected_keys - {"metadata"}} if not isinstance(metadata, dict): raise ArchiveError(f"{path}: metadata is not a JSON object") - validate_metadata(metadata) + metadata = upgrade_metadata(metadata) splits: dict[str, SplitLogits] = {} for name in SPLITS: logits, labels = arrays[f"{name}_logits"], arrays[f"{name}_labels"] @@ -191,6 +279,30 @@ def load_logits(path: Path) -> LogitsArchive: return LogitsArchive(metadata=metadata, splits=splits) +def load_validation_logits(path: Path) -> tuple[dict[str, object], SplitLogits]: + """Metadata and the validation split only; the test arrays are never read. + + For the pilots, which choose hyperparameters and must not touch any + test output (docs/PLAN.md section 4, hyperparameter protocol). An + ``.npz`` is a zip, so reading one member leaves the others unread. + """ + with np.load(path, allow_pickle=False) as npz: + needed = {"metadata", "validation_logits", "validation_labels"} + missing = sorted(needed - set(npz.files)) + if missing: + raise ArchiveError(f"{path}: arrays missing {missing}") + metadata = json.loads(str(npz["metadata"])) + logits, labels = npz["validation_logits"], npz["validation_labels"] + if not isinstance(metadata, dict): + raise ArchiveError(f"{path}: metadata is not a JSON object") + metadata = upgrade_metadata(metadata) + if logits.ndim != 2 or labels.ndim != 1 or logits.shape[0] != labels.shape[0]: + raise ArchiveError(f"{path}: validation logits {logits.shape}, labels {labels.shape}") + split = SplitLogits("validation", logits, labels) + validate_split(split) + return metadata, split + + def read_manifest(manifest_path: Path) -> dict[str, dict[str, object]]: if not manifest_path.exists(): return {} @@ -216,7 +328,7 @@ def record_in_manifest(manifest_path: Path, archive_path: Path) -> dict[str, obj def write_manifest(manifest_path: Path, files: dict[str, dict[str, object]]) -> None: - body = {"format_version": FORMAT_VERSION, "files": dict(sorted(files.items()))} + body = {"format_version": MANIFEST_FORMAT_VERSION, "files": dict(sorted(files.items()))} manifest_path.parent.mkdir(parents=True, exist_ok=True) tmp = manifest_path.with_name(manifest_path.name + ".tmp") tmp.write_text(json.dumps(body, indent=2) + "\n", encoding="utf-8") diff --git a/src/tinyrouter/baselines.py b/src/tinyrouter/baselines.py new file mode 100644 index 0000000..abb36a9 --- /dev/null +++ b/src/tinyrouter/baselines.py @@ -0,0 +1,267 @@ +"""Cheap baselines on every curve point, stored exactly like an encoder run (docs/PLAN.md AC5). + +For each k and seed, both baselines are fit on ``sampling.curve_sample(k, +seed)``, the same rows the encoders train on (the results JSON records the +sample's fingerprint, comparable with an encoder's ``train_sample_sha256``). +Each run writes a logits archive of 151 scores per validation and test +row, an entry in the logits manifest, and ``results/runs/.json`` with +the same metrics block as an encoder, so the step-4 analysis reads all of +them the same way. How the scores are defined: + +- ``majority``: the log of each intent's share of the training sample, + the same 151 numbers for every query (it never looks at the text). Its + argmax is the most frequent intent, which is ``oos`` (ceil(2.5k) rows + against k for every other intent). Summed per agent it is the most + frequent agent, ``finance_agent`` (38 of the 150 intents): that is the + 8-way majority class, so ``accuracy_8_summed`` is the majority-class + accuracy and ``accuracy_8`` (argmax intent, then its agent) always says + oos. Its softmax is the training class distribution, a calibrated + prior; the temperature fit on validation may still move it. +- ``tfidf-centroid``: cosine similarity between the query's TF-IDF vector + and each intent's centroid. The vectorizer (word unigrams and bigrams, + sublinear tf, L2-normalised rows) is fit on the training sample only; + a centroid is the mean of its intent's training vectors, L2-normalised. + The score is ``COSINE_SCALE`` (100) times the cosine, the convention + CLIP uses to turn cosines into logits. The factor changes no argmax; + it only moves the scores to a range where the validation temperature + fit (search range T in [0.05, 100]) has an interior optimum. On the + raw cosines (range [0, 1]) the fitted T hit the lower bound at every + k >= 5 (checked on seed 42 before this was written). A query with no + known term scores 0 everywhere and its argmax is intent 0. +""" + +from __future__ import annotations + +import argparse +import json +import os +from collections.abc import Callable +from pathlib import Path + +import numpy as np + +from tinyrouter.archive import ( + FORMAT_VERSION, + git_state, + load_logits, + pinned_dataset_sha256, + record_in_manifest, + save_logits, + utc_now, +) +from tinyrouter.calibrate import SplitLogits +from tinyrouter.data import DATASET_REVISION, Split, load_split +from tinyrouter.evaluate import RunPaths, score +from tinyrouter.labels import load_label_space +from tinyrouter.protocol import SEEDS +from tinyrouter.runs import Log, archive_intact, read_record +from tinyrouter.sampling import CURVE_KS, curve_sample, sample_fingerprint + +# Bump when a definition below changes; a stored run with another version is redone. +BASELINE_VERSION = "1" +TFIDF_PARAMS: dict[str, object] = {"ngram_range": (1, 2), "sublinear_tf": True} +COSINE_SCALE = 100.0 +DEFINITIONS = { + "majority": "log share of each intent in the training sample, same for every query", + "tfidf-centroid": "100 * cosine(query tf-idf, L2-normalised intent centroid); vectorizer " + "fit on the training sample, word 1-2 grams, sublinear tf", +} +BASELINES = tuple(DEFINITIONS) +SampleFn = Callable[[int, int], Split] +EvalFn = Callable[[str], Split] + + +def majority_scores(train: Split, rows: int) -> np.ndarray: + counts = np.bincount(train.intents, minlength=load_label_space().num_intents) + if (counts == 0).any(): + raise ValueError(f"intents {np.flatnonzero(counts == 0).tolist()} have no training rows") + prior = np.log(counts / counts.sum()) + return np.tile(prior, (rows, 1)).astype(np.float32) + + +class TfidfCentroid: + """Nearest-centroid scorer over TF-IDF vectors, fit on the training sample only.""" + + def __init__(self, train: Split) -> None: + from sklearn.feature_extraction.text import TfidfVectorizer + from sklearn.preprocessing import normalize + + self.vectorizer = TfidfVectorizer(**TFIDF_PARAMS) # type: ignore[arg-type] + vectors = self.vectorizer.fit_transform(train.texts) + num_intents = load_label_space().num_intents + centroids = np.zeros((num_intents, vectors.shape[1])) + for intent in range(num_intents): + rows = np.flatnonzero(train.intents == intent) + if len(rows) == 0: + raise ValueError(f"intent {intent} has no training rows") + centroids[intent] = np.asarray(vectors[rows].mean(axis=0)).ravel() + self.centroids = normalize(centroids) + + def scores(self, texts: tuple[str, ...]) -> np.ndarray: + cosine = self.vectorizer.transform(texts) @ self.centroids.T + return (COSINE_SCALE * np.asarray(cosine)).astype(np.float32) + + +def baseline_logits(name: str, train: Split, evals: dict[str, Split]) -> dict[str, SplitLogits]: + if name == "majority": + make = lambda split: majority_scores(train, len(split)) # noqa: E731 + elif name == "tfidf-centroid": + model = TfidfCentroid(train) + make = lambda split: model.scores(split.texts) # noqa: E731 + else: + raise ValueError(f"unknown baseline '{name}'; expected one of {BASELINES}") + return {n: SplitLogits(n, make(split), split.intents) for n, split in evals.items()} # type: ignore[arg-type] + + +def run_name(name: str, k: int, seed: int) -> str: + return f"{name}-k{k}-seed{seed}" + + +def definition(name: str) -> dict[str, object]: + params = {k: list(v) if isinstance(v, tuple) else v for k, v in TFIDF_PARAMS.items()} + return { + "name": name, + "version": BASELINE_VERSION, + "scores": DEFINITIONS[name], + "params": params if name == "tfidf-centroid" else {}, + } + + +def training_block(train: Split, k: int, seed: int) -> dict[str, object]: + oos = load_label_space().oos_intent_id + return { + "seed": seed, + "k_shot": k, + "train_rows": len(train), + "oos_train_rows": int((train.intents == oos).sum()), + "train_sample_sha256": sample_fingerprint(train), + } + + +def metadata(name: str, k: int, seed: int, training: dict[str, object]) -> dict[str, object]: + commit, dirty = git_state() + return { + "format_version": FORMAT_VERSION, + "run_name": run_name(name, k, seed), + "model_name": f"baseline/{name}", + "model_revision": f"v{BASELINE_VERSION}", + "seed": seed, + "per_intent": None, + "k_shot": k, + "train_rows": training["train_rows"], + "oos_train_rows": training["oos_train_rows"], + "eval_per_intent": None, + "dataset_revision": DATASET_REVISION, + "dataset_sha256": pinned_dataset_sha256(), + "label_space_sha256": load_label_space().sha256, + "git_commit": commit, + "git_dirty": dirty, + "created_at": utc_now(), + } + + +def is_done(paths: RunPaths, name: str, training: dict[str, object], seed: int) -> bool: + record = read_record(paths) + if record is None or record.get("baseline") != definition(name): + return False + if record.get("training") != training: + return False + return archive_intact(paths, record, seed) + + +def run_baseline( + name: str, k: int, seed: int, train: Split, evals: dict[str, Split], results_root: Path +) -> dict[str, object]: + """Fit, score, archive and record one baseline run; skip if an identical one is intact.""" + paths = RunPaths.named(results_root, run_name(name, k, seed)) + training = training_block(train, k, seed) + if is_done(paths, name, training, seed): + record = read_record(paths) + assert record is not None + return record + paths.results_json.unlink(missing_ok=True) + splits = baseline_logits(name, train, evals) + save_logits(paths.logits, splits, metadata(name, k, seed, training)) + entry = record_in_manifest(paths.manifest, paths.logits) + archived = load_logits(paths.logits) + record = { + "run_name": run_name(name, k, seed), + "baseline": definition(name), + "environment": environment(), + "training": training, + "logits": {"file": paths.logits.name, "sha256": entry["sha256"], "bytes": entry["bytes"]}, + "metrics": score(archived.validation, archived.test), + } + paths.results_json.parent.mkdir(parents=True, exist_ok=True) + paths.results_json.write_text(json.dumps(record, indent=2) + "\n", encoding="utf-8") + return record + + +def environment() -> dict[str, object]: + import platform + + import sklearn + + return { + "python": platform.python_version(), + "numpy": np.__version__, + "scikit-learn": sklearn.__version__, + "dataset_revision": DATASET_REVISION, + } + + +def run_all( + results_root: Path, + sample_fn: SampleFn = curve_sample, + eval_fn: EvalFn = load_split, # type: ignore[assignment] + log: Log = print, +) -> dict[str, object]: + """Every baseline at every (k, seed); writes results/curves/baselines.json.""" + evals = {name: eval_fn(name) for name in ("validation", "test")} + entries = [] + for k in CURVE_KS: + for seed in SEEDS: + train = sample_fn(k, seed) + for name in BASELINES: + record = run_baseline(name, k, seed, train, evals, results_root) + entries.append(index_entry(name, k, seed, record)) + log(f"baselines k={k} seed={seed}: done") + body = {"baselines": {n: definition(n) for n in BASELINES}, "points": entries} + out = results_root / "curves" / "baselines.json" + write_index(out, body) + return body + + +def index_entry(name: str, k: int, seed: int, record: dict[str, object]) -> dict[str, object]: + training, logits = record["training"], record["logits"] + assert isinstance(training, dict) and isinstance(logits, dict) + return { + "baseline": name, + "k": k, + "seed": seed, + "run_name": record["run_name"], + "logits_file": logits["file"], + "logits_sha256": logits["sha256"], + "train_rows": training["train_rows"], + "oos_train_rows": training["oos_train_rows"], + "train_sample_sha256": training["train_sample_sha256"], + } + + +def write_index(path: Path, body: dict[str, object]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + tmp = path.with_name(path.name + ".tmp") + tmp.write_text(json.dumps(body, indent=2) + "\n", encoding="utf-8") + os.replace(tmp, path) + + +def main(argv: list[str] | None = None) -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--results-root", default="results") + args = parser.parse_args(argv) + body = run_all(Path(args.results_root)) + print(f"wrote {len(body['points'])} baseline runs") # type: ignore[arg-type] + + +if __name__ == "__main__": + main() diff --git a/src/tinyrouter/config.py b/src/tinyrouter/config.py index dd56bc1..5c45aff 100644 --- a/src/tinyrouter/config.py +++ b/src/tinyrouter/config.py @@ -13,6 +13,16 @@ # Where output goes, not what it is: two configs differing only here produce the same run. LOCATION_FIELDS = frozenset({"checkpoint_root", "results_root"}) +# Fields added after results were first recorded (AC2, format of 2026-09-23), +# with the default that reproduces the behaviour of the code before they +# existed. A recorded config that lacks one of these keys is read as having +# this default; any other missing key still means "not the same run". +ADDED_FIELD_DEFAULTS: dict[str, object] = { + "k_shot": None, + "oos_train": None, + "min_train_steps": None, +} + @dataclass(frozen=True) class RunConfig: @@ -20,12 +30,20 @@ class RunConfig: model_revision: str seed: int = 42 per_intent: int | None = None + # Learning-curve sample (sampling.py): k rows per intent, oos from the + # hardcoded table unless ``oos_train`` overrides it (OOS ablation). + k_shot: int | None = None + oos_train: int | None = None max_length: int = 64 - learning_rate: float = 5e-5 + # None means "not chosen yet" (a pilot decides it); training refuses it. + learning_rate: float | None = 5e-5 weight_decay: float = 0.01 warmup_ratio: float = 0.1 num_train_epochs: float = 5.0 max_steps: int = -1 + # Train for max(min_train_steps, the steps num_train_epochs gives). + # None trains by epochs only. See train.plan_steps. + min_train_steps: int | None = None train_batch_size: int = 32 eval_batch_size: int = 128 device: Device = "auto" @@ -42,10 +60,26 @@ class RunConfig: def __post_init__(self) -> None: if not 0.0 <= self.warmup_ratio < 1.0: raise ValueError(f"warmup_ratio must be in [0, 1), got {self.warmup_ratio}") + if self.k_shot is not None and self.per_intent is not None: + raise ValueError("set k_shot (curve sample) or per_intent (per-intent cap), not both") + if self.oos_train is not None and self.k_shot is None: + raise ValueError("oos_train overrides the k-shot oos count; it needs k_shot") + if self.min_train_steps is not None: + if self.min_train_steps < 1: + raise ValueError(f"min_train_steps must be >= 1, got {self.min_train_steps}") + if self.max_steps > 0: + raise ValueError("min_train_steps and max_steps > 0 contradict each other") @property def run_name(self) -> str: - size = "full" if self.per_intent is None else f"k{self.per_intent}" + if self.k_shot is not None: + size = f"k{self.k_shot}" + if self.oos_train is not None: + size += f"-oos{self.oos_train}" + elif self.per_intent is not None: + size = f"cap{self.per_intent}" + else: + size = "full" short = self.model_name.rstrip("/").split("/")[-1] return f"{short}-{size}-seed{self.seed}" @@ -60,10 +94,16 @@ def identity(self) -> dict[str, object]: return {k: v for k, v in asdict(self).items() if k not in LOCATION_FIELDS} def matches(self, recorded: object) -> bool: - """Whether a config dict saved with earlier output describes this same run.""" + """Whether a config dict saved with earlier output describes this same run. + + Exact field-by-field equality, except that keys in + ``ADDED_FIELD_DEFAULTS`` missing from an older record count as + their default. + """ if not isinstance(recorded, dict): return False - return {k: v for k, v in recorded.items() if k not in LOCATION_FIELDS} == self.identity() + filled = {**ADDED_FIELD_DEFAULTS, **recorded} + return {k: v for k, v in filled.items() if k not in LOCATION_FIELDS} == self.identity() def with_seed(self, seed: int) -> RunConfig: return replace(self, seed=seed) diff --git a/src/tinyrouter/curves.py b/src/tinyrouter/curves.py new file mode 100644 index 0000000..16b4ca0 --- /dev/null +++ b/src/tinyrouter/curves.py @@ -0,0 +1,149 @@ +"""Learning curves (RQ1) and the OOS ablation (RQ2), docs/PLAN.md sections 3 to 5. + +``make curve MODEL=bert|modernbert``: the baselines (baselines.py), then +6 values of k x 3 seeds with the +learning rate and S_min from ``configs/curve.yaml``. ``make oos-ablation``: +ModernBERT at k=100 with 0 OOS training rows x 3 seeds. Both resume with +the rules in runs.py (a run is skipped only when its results JSON came +from this config and its archive's SHA-256 agrees in three places), and +no curve weights are kept: each run's weights are deleted once its logits +are archived. + +BERT at k=100 can be the AC2 run. ``protocol.ac2_donor`` proposes the AC2 +run of the same seed, and ``runs.reusable_equivalent`` accepts it only if +training with the curve config would have been the same run (same fields, +same rows, same step count, and the AC2 record's own row and step counts +agree). If the pilot picked a learning rate other than 5e-5, or anything +else differs, the point is trained like any other. + +Each command writes an index, ``results/curves/.json``, with one +entry per (k, seed): the run that holds its logits, the archive's SHA-256, +the planned steps and which rule decided them (``epochs`` or +``min_train_steps``), the steps actually run, and ``reused_from`` when the +point is an AC2 run. +""" + +from __future__ import annotations + +import argparse +import json +import os +from pathlib import Path + +from tinyrouter.baselines import run_all as run_baselines +from tinyrouter.config import RunConfig +from tinyrouter.evaluate import RunPaths +from tinyrouter.protocol import MODELS, CurveProtocol, ac2_donor, load_protocol +from tinyrouter.runs import ( + EvaluateFn, + Log, + TrainFn, + default_evaluate, + default_train, + read_record, + reusable_equivalent, + run_one, +) +from tinyrouter.steps import plan_steps + + +class CurveError(RuntimeError): + """A finished curve point does not look like the run it was supposed to be.""" + + +def index_entry(config: RunConfig, record: dict, reused_from: str | None) -> dict[str, object]: + training, logits = record["training"], record["logits"] + plan = plan_steps(config, training["train_rows"]) + if training["global_step"] != plan.planned_steps: + raise CurveError( + f"{record['run_name']}: ran {training['global_step']} steps, " + f"{config.run_name} plans {plan.planned_steps}" + ) + if reused_from is None and training.get("step_plan") != plan.as_dict(): + raise CurveError(f"{record['run_name']}: recorded step plan differs from this config's") + return { + "k": config.k_shot, + "seed": config.seed, + "oos_train": config.oos_train, + "run_name": record["run_name"], + "logits_file": logits["file"], + "logits_sha256": logits["sha256"], + "train_rows": training["train_rows"], + "oos_train_rows": training["oos_train_rows"], + "train_sample_sha256": training.get("train_sample_sha256"), + "planned_steps": plan.planned_steps, + "epoch_steps": plan.epoch_steps, + "decided_by": plan.decided_by, + "global_step": training["global_step"], + "reused_from": reused_from, + } + + +def run_point( + config: RunConfig, protocol: CurveProtocol, train_fn: TrainFn, evaluate_fn: EvaluateFn, log: Log +) -> dict[str, object]: + donor = ac2_donor(protocol, config) + if donor is not None: + record = reusable_equivalent(config, donor) + if record is not None: + log(f"{config.run_name}: reusing {donor.run_name} (equivalent run, archive intact)") + return index_entry(config, record, donor.run_name) + log(f"{config.run_name}: {donor.run_name} is not an equivalent intact run; training") + run_one(config, train_fn, evaluate_fn, log, keep_weights=False) + record = read_record(RunPaths.of(config)) + if record is None: + raise CurveError(f"{config.run_name}: no results JSON after evaluation") + return index_entry(config, record, None) + + +def run_curve( + name: str, + configs: list[RunConfig], + protocol: CurveProtocol, + train_fn: TrainFn = default_train, + evaluate_fn: EvaluateFn = default_evaluate, + log: Log = print, +) -> dict[str, object]: + """Run or resume every config, then write results/curves/.json.""" + entries = [run_point(c, protocol, train_fn, evaluate_fn, log) for c in configs] + shas = [e["logits_sha256"] for e in entries] + if len(set(shas)) != len(shas): + raise CurveError(f"{name}: two points share one logits archive") + first = configs[0] + body = { + "name": name, + "model_name": first.model_name, + "model_revision": first.model_revision, + "learning_rate": first.learning_rate, + "min_train_steps": first.min_train_steps, + "points": entries, + } + out = Path(first.results_root) / "curves" / f"{name}.json" + out.parent.mkdir(parents=True, exist_ok=True) + tmp = out.with_name(out.name + ".tmp") + tmp.write_text(json.dumps(body, indent=2) + "\n", encoding="utf-8") + os.replace(tmp, out) + log(f"wrote {out} ({len(entries)} points)") + return body + + +def main(argv: list[str] | None = None) -> None: + parser = argparse.ArgumentParser(description=__doc__) + which = parser.add_mutually_exclusive_group(required=True) + which.add_argument("--model", choices=MODELS) + which.add_argument("--ablation", action="store_true", help="ModernBERT, k=100, 0 OOS rows") + parser.add_argument("--protocol", default="configs/curve.yaml") + args = parser.parse_args(argv) + protocol = load_protocol(args.protocol) + # Build every config first: a protocol still missing a pilot value stops here, + # before anything is computed or written. + if args.ablation: + run_curve("oos-ablation", protocol.ablation_configs(), protocol) + return + configs = protocol.curve_configs(args.model) + run_baselines(Path(configs[0].results_root)) + run_curve(args.model, configs, protocol) + + +if __name__ == "__main__": + main() diff --git a/src/tinyrouter/evaluate.py b/src/tinyrouter/evaluate.py index 377e603..cbbab73 100644 --- a/src/tinyrouter/evaluate.py +++ b/src/tinyrouter/evaluate.py @@ -24,9 +24,11 @@ import numpy as np from tinyrouter.archive import ( + FORMAT_VERSION, MANIFEST_NAME, git_state, load_logits, + pinned_dataset_sha256, record_in_manifest, save_logits, utc_now, @@ -47,10 +49,14 @@ class RunPaths: @classmethod def of(cls, config: RunConfig) -> RunPaths: - root = Path(config.results_root) + return cls.named(config.results_root, config.run_name) + + @classmethod + def named(cls, results_root: str | Path, run_name: str) -> RunPaths: + root = Path(results_root) return cls( - results_json=root / "runs" / f"{config.run_name}.json", - logits=root / "logits" / f"{config.run_name}.npz", + results_json=root / "runs" / f"{run_name}.json", + logits=root / "logits" / f"{run_name}.npz", manifest=root / MANIFEST_NAME, ) @@ -165,16 +171,19 @@ def read_training_summary(model_dir: Path) -> dict[str, object]: def archive_metadata(config: RunConfig, training: dict[str, object]) -> dict[str, object]: commit, dirty = git_state() return { - "format_version": 1, + "format_version": FORMAT_VERSION, "run_name": config.run_name, "model_name": config.model_name, "model_revision": config.model_revision, "seed": config.seed, "per_intent": config.per_intent, + "k_shot": config.k_shot, "train_rows": training["train_rows"], "oos_train_rows": training["oos_train_rows"], "eval_per_intent": config.eval_per_intent, "dataset_revision": DATASET_REVISION, + # load_split verified every file against these before anything read it. + "dataset_sha256": pinned_dataset_sha256(), "label_space_sha256": load_label_space().sha256, "git_commit": commit, "git_dirty": dirty, diff --git a/src/tinyrouter/metrics.py b/src/tinyrouter/metrics.py index 4a42eca..645d0a4 100644 --- a/src/tinyrouter/metrics.py +++ b/src/tinyrouter/metrics.py @@ -36,6 +36,27 @@ def oos_recall(pred: np.ndarray, gold: np.ndarray, oos_id: int) -> float: return float(np.mean(pred[mask] == oos_id)) +def wilson_interval(successes: int, n: int, z: float = 1.96) -> tuple[float, float]: + """Wilson score interval for a binomial proportion ``successes / n``. + + Used for OOS recall, where the test split has 1,000 OOS rows and a + rate near 0 or 1 makes the normal-approximation interval fall outside + [0, 1]. The Wilson interval stays inside [0, 1] and is not degenerate + at 0 or n successes. ``z = 1.96`` gives a two-sided 95% interval. + """ + if n < 1: + raise ValueError(f"n must be >= 1, got {n}") + if not 0 <= successes <= n: + raise ValueError(f"successes must be in [0, {n}], got {successes}") + if z <= 0: + raise ValueError(f"z must be positive, got {z}") + p = successes / n + z2 = z * z + centre = (p + z2 / (2 * n)) / (1 + z2 / n) + half = z * np.sqrt(p * (1 - p) / n + z2 / (4 * n * n)) / (1 + z2 / n) + return float(max(0.0, centre - half)), float(min(1.0, centre + half)) + + def expected_calibration_error(probs: np.ndarray, gold: np.ndarray, n_bins: int = 15) -> float: """Top-label ECE with equal-width confidence bins over (0, 1]. diff --git a/src/tinyrouter/pilots.py b/src/tinyrouter/pilots.py new file mode 100644 index 0000000..5fdb205 --- /dev/null +++ b/src/tinyrouter/pilots.py @@ -0,0 +1,329 @@ +"""Validation-only pilots that fix the curve's learning rates and S_min (docs/PLAN.md section 4). + +``make pilot-lr``: each encoder at k=100, seed 42, learning rate in +{1e-5, 2e-5, 5e-5}. Chosen by validation 150-way in-scope accuracy, then +validation OOS recall, then the smaller learning rate. BERT at 5e-5 is the +AC2 seed-42 run, reused when ``runs.equivalent_run`` confirms it; only its +validation logits are read. + +``make pilot-steps``: both encoders at k=5, seed 42, with their chosen +learning rates, S_min in {100, 200, 400}. Chosen by the mean of the two +encoders' validation in-scope accuracy, then the smaller S_min. + +Neither pilot scores, loads or writes anything from the test split: +``predict_validation`` and ``validation_scores`` raise ``LeakageError`` on +any other split, the reused AC2 archive is read through +``load_validation_logits``, and nothing here calls ``evaluate()`` (which +scores test). Outputs go to ``results/pilots/{lr,steps}.json``. The +selected values are printed, not written into ``configs/curve.yaml``; +copying them there is a deliberate manual step. + +Ties are compared on integer counts of correct rows, not on floats, and +every point is scored on the same 3,100 validation rows. +""" + +from __future__ import annotations + +import argparse +import json +import os +import shutil +from collections.abc import Callable +from dataclasses import dataclass, replace +from pathlib import Path + +from tinyrouter.archive import load_validation_logits +from tinyrouter.calibrate import LeakageError, SplitLogits +from tinyrouter.config import RunConfig +from tinyrouter.data import Split +from tinyrouter.evaluate import RunPaths, read_training_summary +from tinyrouter.labels import load_label_space +from tinyrouter.protocol import ( + LR_GRID, + MODELS, + PILOT_LR_K, + PILOT_SEED, + PILOT_STEPS_K, + S_MIN_GRID, + CurveProtocol, + ac2_donor, + load_protocol, +) +from tinyrouter.runs import ( + Log, + TrainFn, + default_train, + reusable_equivalent, + reusable_weights, + run_dir, +) + +ValidationFn = Callable[[RunConfig, Path], SplitLogits] +PILOT_DIR = "pilots" +# Training-summary fields copied into a pilot point (no test data exists in a summary). +TRAINING_FIELDS = ( + "train_rows", + "oos_train_rows", + "k_shot", + "train_sample_sha256", + "step_plan", + "global_step", + "train_wall_seconds", +) + + +def require_validation(split_name: str) -> None: + if split_name != "validation": + raise LeakageError( + f"pilots choose hyperparameters and may only see validation, got '{split_name}'" + ) + + +def predict_validation(model_dir: Path, split: Split, config: RunConfig) -> SplitLogits: + """Logits of ``split``, which must be validation.""" + from tinyrouter.evaluate import predict_logits + + require_validation(split.name) + return predict_logits(model_dir, split, config) + + +def validation_scores(val: SplitLogits) -> dict[str, float | int]: + """In-scope accuracy and OOS recall on validation, with the integer counts ties use.""" + require_validation(val.split) + oos = load_label_space().oos_intent_id + pred, gold = val.logits.argmax(axis=1), val.labels + in_scope, is_oos = gold != oos, gold == oos + in_correct = int((pred[in_scope] == gold[in_scope]).sum()) + oos_correct = int((pred[is_oos] == oos).sum()) + n_in, n_oos = int(in_scope.sum()), int(is_oos.sum()) + return { + "in_scope_correct": in_correct, + "in_scope_n": n_in, + "in_scope_accuracy_150": in_correct / n_in, + "oos_correct": oos_correct, + "oos_n": n_oos, + "oos_recall_151": oos_correct / n_oos, + "n": int(len(gold)), + } + + +def default_validation(config: RunConfig, model_dir: Path) -> SplitLogits: + from tinyrouter.data import load_split + + return predict_validation(model_dir, load_split("validation"), config) + + +@dataclass(frozen=True) +class PilotPoint: + model: str + value: float | int + config: RunConfig + + +def pilot_location(config: RunConfig, tag: str) -> RunConfig: + return replace(config, checkpoint_root=str(Path(config.checkpoint_root) / PILOT_DIR / tag)) + + +def lr_points(protocol: CurveProtocol) -> list[PilotPoint]: + """k=100 already exceeds every S_min (2,385 epoch steps), so S_min is left unset here.""" + points = [] + for model in MODELS: + for lr in LR_GRID: + config = replace( + protocol.base(model), + learning_rate=lr, + k_shot=PILOT_LR_K, + oos_train=None, + min_train_steps=None, + seed=PILOT_SEED, + ) + points.append(PilotPoint(model, lr, pilot_location(config, f"lr{lr:g}"))) + return points + + +def steps_points(protocol: CurveProtocol) -> list[PilotPoint]: + points = [] + for s_min in S_MIN_GRID: + for model in MODELS: + config = replace( + protocol.base(model), + learning_rate=protocol.learning_rate(model), + k_shot=PILOT_STEPS_K, + oos_train=None, + min_train_steps=s_min, + seed=PILOT_SEED, + ) + points.append(PilotPoint(model, s_min, pilot_location(config, f"smin{s_min}"))) + return points + + +def point_entry(point: PilotPoint, training: dict, scores: dict, source: str) -> dict: + return { + "model": point.model, + "value": point.value, + "run_name": point.config.run_name, + "config": point.config.identity(), + "source": source, + "training": {key: training.get(key) for key in TRAINING_FIELDS}, + "validation": scores, + } + + +def reused_entry(point: PilotPoint, protocol: CurveProtocol) -> dict | None: + """The AC2 run as this point, if it is an equivalent run; validation logits only.""" + donor = ac2_donor(protocol, point.config) + record = None if donor is None else reusable_equivalent(point.config, donor) + if donor is None or record is None: + return None + _, val = load_validation_logits(RunPaths.of(donor).logits) + training = record["training"] + assert isinstance(training, dict) + source = f"reused {donor.run_name} (equivalent run; validation logits from its archive)" + return point_entry(point, training, validation_scores(val), source) + + +def trained_entry( + point: PilotPoint, train_fn: TrainFn, validation_fn: ValidationFn, log: Log +) -> dict: + if reusable_weights(point.config): + final = run_dir(point.config) / "final" + log(f"{point.model} {point.value}: reusing weights trained with this config") + else: + if run_dir(point.config).exists(): + shutil.rmtree(run_dir(point.config)) + log(f"{point.model} {point.value}: training {point.config.run_name}") + final = train_fn(point.config) + training = read_training_summary(final) + scores = validation_scores(validation_fn(point.config, final)) + entry = point_entry(point, training, scores, "trained") + shutil.rmtree(run_dir(point.config)) + return entry + + +def select_lr(entries: list[dict]) -> dict[str, float]: + """Per model: most in-scope correct, then most OOS correct, then the smaller rate.""" + chosen = {} + for model in MODELS: + mine = [e for e in entries if e["model"] == model] + same_rows(mine) + best = max( + mine, + key=lambda e: ( + e["validation"]["in_scope_correct"], + e["validation"]["oos_correct"], + -e["value"], + ), + ) + chosen[model] = best["value"] + return chosen + + +def select_steps(entries: list[dict]) -> int: + """Largest total in-scope correct over both encoders (same as the mean), then smaller S_min.""" + same_rows(entries) + totals = { + s_min: sum(e["validation"]["in_scope_correct"] for e in entries if e["value"] == s_min) + for s_min in S_MIN_GRID + } + return max(S_MIN_GRID, key=lambda s: (totals[s], -s)) + + +def same_rows(entries: list[dict]) -> None: + sizes = {(e["validation"]["in_scope_n"], e["validation"]["oos_n"]) for e in entries} + if len(sizes) != 1: + raise ValueError(f"pilot points were scored on different validation sets: {sizes}") + + +def load_previous(out: Path) -> list[dict]: + if not out.exists(): + return [] + return list(json.loads(out.read_text(encoding="utf-8")).get("points", [])) + + +def write_json(path: Path, body: dict) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + tmp = path.with_name(path.name + ".tmp") + tmp.write_text(json.dumps(body, indent=2) + "\n", encoding="utf-8") + os.replace(tmp, path) + + +def run_pilot( + kind: str, + protocol: CurveProtocol, + out: Path, + train_fn: TrainFn = default_train, + validation_fn: ValidationFn = default_validation, + log: Log = print, +) -> dict: + """Run (or resume) one pilot and write its JSON; returns the body with the selection.""" + points = lr_points(protocol) if kind == "lr" else steps_points(protocol) + previous = load_previous(out) + entries: list[dict] = [] + for point in points: + done = next( + ( + e + for e in previous + if e["model"] == point.model + and e["value"] == point.value + and point.config.matches(e["config"]) + ), + None, + ) + if done is not None: + log(f"{point.model} {point.value}: already scored with this config, skipping") + entries.append(done) + continue + entry = reused_entry(point, protocol) if kind == "lr" else None + entries.append(entry or trained_entry(point, train_fn, validation_fn, log)) + write_json(out, pilot_body(kind, entries, selected=None)) + selected = select_lr(entries) if kind == "lr" else select_steps(entries) + body = pilot_body(kind, entries, selected) + write_json(out, body) + return body + + +def pilot_body(kind: str, entries: list[dict], selected: object) -> dict: + if kind == "lr": + setup = {"k": PILOT_LR_K, "seed": PILOT_SEED, "grid": list(LR_GRID)} + rule = "per model: max validation in_scope_correct, then oos_correct, then smaller lr" + else: + setup = {"k": PILOT_STEPS_K, "seed": PILOT_SEED, "grid": list(S_MIN_GRID)} + rule = "sum of both models' validation in_scope_correct, then smaller S_min" + return { + "pilot": "learning_rate" if kind == "lr" else "min_train_steps", + "split": "validation", + **setup, + "rule": rule, + "points": entries, + "selected": selected, + "note": "not applied automatically; copy into configs/curve.yaml by hand", + } + + +def print_summary(body: dict) -> None: + for e in body["points"]: + v = e["validation"] + print( + f"{e['model']:<11} {e['value']!s:<7} in-scope {v['in_scope_accuracy_150']:.4f} " + f"({v['in_scope_correct']}/{v['in_scope_n']}) OOS recall {v['oos_recall_151']:.4f} " + f"steps {e['training']['global_step']} {e['source']}" + ) + print(f"selected {body['pilot']}: {body['selected']} (validation only)") + + +def main(argv: list[str] | None = None) -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("kind", choices=("lr", "steps")) + parser.add_argument("--protocol", default="configs/curve.yaml") + parser.add_argument("--results-root", default="results") + args = parser.parse_args(argv) + protocol = load_protocol(args.protocol) + out = Path(args.results_root) / "pilots" / f"{args.kind}.json" + body = run_pilot(args.kind, protocol, out) + print_summary(body) + print(f"wrote {out}") + + +if __name__ == "__main__": + main() diff --git a/src/tinyrouter/protocol.py b/src/tinyrouter/protocol.py new file mode 100644 index 0000000..1b53b28 --- /dev/null +++ b/src/tinyrouter/protocol.py @@ -0,0 +1,141 @@ +"""Learning-curve hyperparameter protocol (docs/PLAN.md section 4, hyperparameter protocol). + +``configs/curve.yaml`` holds the two values the pilots choose: one +``min_train_steps`` (S_min) shared by every k and both encoders, and one +learning rate per encoder shared by its whole curve. Both are written by +hand from ``results/pilots/*.json``; a null means that pilot has not been +run, and the curve refuses to start. Values outside the protocol's grids +are refused too, so a typo cannot quietly become a new setting. + +Every curve run is the model's base config (``configs/.yaml``, +which also fixes max_length, batch size, epochs, warmup and weight decay) +with the chosen learning rate and S_min, plus a k and a seed. +""" + +from __future__ import annotations + +from dataclasses import dataclass, replace +from pathlib import Path + +import yaml + +from tinyrouter.config import RunConfig, load_config +from tinyrouter.sampling import CURVE_KS, FULL_K + +MODELS = ("bert", "modernbert") +SEEDS = (42, 43, 44) +LR_GRID = (1e-5, 2e-5, 5e-5) +S_MIN_GRID = (100, 200, 400) +PILOT_SEED = 42 +PILOT_LR_K = 100 +PILOT_STEPS_K = 5 +ABLATION_MODEL = "modernbert" +ABLATION_K = 100 +ABLATION_OOS_ROWS = 0 +DEFAULT_PROTOCOL = Path("configs/curve.yaml") + + +class ProtocolError(ValueError): + """configs/curve.yaml is malformed, off the protocol's grids, or not filled in yet.""" + + +@dataclass(frozen=True) +class CurveProtocol: + min_train_steps: int | None + learning_rates: dict[str, float | None] + base_configs: dict[str, RunConfig] + base_paths: dict[str, str] + + def base(self, model: str) -> RunConfig: + """The model's config file as is (for BERT this is the AC2 config).""" + if model not in MODELS: + raise ProtocolError(f"unknown model '{model}'; expected one of {MODELS}") + return self.base_configs[model] + + def learning_rate(self, model: str) -> float: + lr = self.learning_rates[model] + if lr is None: + raise ProtocolError( + f"configs/curve.yaml: learning_rate for {model} is null; run `make pilot-lr` " + "and copy the selected value in" + ) + return lr + + def s_min(self) -> int: + if self.min_train_steps is None: + raise ProtocolError( + "configs/curve.yaml: min_train_steps is null; run `make pilot-steps` and copy " + "the selected value in" + ) + return self.min_train_steps + + def curve_config( + self, model: str, k: int, seed: int, oos_train: int | None = None + ) -> RunConfig: + """One curve (or ablation) run; needs both pilots decided.""" + if k not in CURVE_KS: + raise ProtocolError(f"k={k} is not a curve point {CURVE_KS}") + return replace( + self.base(model), + learning_rate=self.learning_rate(model), + min_train_steps=self.s_min(), + k_shot=k, + oos_train=oos_train, + seed=seed, + ) + + def curve_configs(self, model: str) -> list[RunConfig]: + return [self.curve_config(model, k, seed) for k in CURVE_KS for seed in SEEDS] + + def ablation_configs(self) -> list[RunConfig]: + return [ + self.curve_config(ABLATION_MODEL, ABLATION_K, seed, oos_train=ABLATION_OOS_ROWS) + for seed in SEEDS + ] + + +def ac2_donor(protocol: CurveProtocol, config: RunConfig) -> RunConfig | None: + """The AC2 run that could stand in for ``config``: BERT's own config file at this seed. + + Only a candidate. Whether it really is the same run is decided by + ``runs.reusable_equivalent``, which compares every field, the rows and + the step count, and checks the archive's SHA-256 in three places. + """ + bert = protocol.base("bert") + if config.model_name != bert.model_name or config.k_shot != FULL_K: + return None + if config.oos_train is not None: + return None + return bert.with_seed(config.seed) + + +def _on_grid(value: object, grid: tuple, name: str) -> None: + if value is None: + return + if isinstance(value, bool) or not isinstance(value, int | float) or value not in grid: + raise ProtocolError(f"{name} = {value!r} is not one of the protocol's values {grid}") + + +def load_protocol(path: str | Path = DEFAULT_PROTOCOL) -> CurveProtocol: + raw = yaml.safe_load(Path(path).read_text(encoding="utf-8")) + if not isinstance(raw, dict) or set(raw) != {"min_train_steps", "models"}: + raise ProtocolError(f"{path}: top level must have exactly min_train_steps and models") + models = raw["models"] + if not isinstance(models, dict) or set(models) != set(MODELS): + raise ProtocolError(f"{path}: models must be exactly {MODELS}") + _on_grid(raw["min_train_steps"], S_MIN_GRID, "min_train_steps") + if isinstance(raw["min_train_steps"], float): + raise ProtocolError(f"min_train_steps must be an integer, got {raw['min_train_steps']!r}") + lrs: dict[str, float | None] = {} + bases: dict[str, RunConfig] = {} + paths: dict[str, str] = {} + for model in MODELS: + entry = models[model] + if not isinstance(entry, dict) or set(entry) != {"config", "learning_rate"}: + raise ProtocolError(f"{path}: models.{model} needs exactly config and learning_rate") + _on_grid(entry["learning_rate"], LR_GRID, f"models.{model}.learning_rate") + lr = entry["learning_rate"] + lrs[model] = None if lr is None else float(lr) + paths[model] = str(entry["config"]) + bases[model] = load_config(entry["config"]) + return CurveProtocol(raw["min_train_steps"], lrs, bases, paths) diff --git a/src/tinyrouter/runs.py b/src/tinyrouter/runs.py new file mode 100644 index 0000000..e8a6288 --- /dev/null +++ b/src/tinyrouter/runs.py @@ -0,0 +1,225 @@ +"""Run lifecycle shared by AC2, the learning curves and the pilots. + +A run is done only when its results JSON was produced by the current +config (output paths aside) and its logits archive has the same SHA-256 in +the results JSON, in the manifest and on disk. Anything else is cleared and +redone. Trained weights are reused only when their ``train_summary.json`` +records the current config. + +``equivalent_run`` answers a narrower question for reusing an existing run +under a different config: would training with ``config`` have run the +same code on the same rows with the same hyperparameters? See its +docstring for the two rewrites it accepts; everything else must be equal. +""" + +from __future__ import annotations + +import json +import shutil +from collections.abc import Callable +from dataclasses import fields +from pathlib import Path + +from tinyrouter.archive import ArchiveError, load_logits, read_manifest, remove_from_manifest +from tinyrouter.config import ADDED_FIELD_DEFAULTS, LOCATION_FIELDS, RunConfig +from tinyrouter.data import sha256_of +from tinyrouter.evaluate import RunPaths +from tinyrouter.sampling import FULL_K, OOS_TRAIN_ROWS, full_train_rows, planned_rows +from tinyrouter.steps import plan_steps + +TrainFn = Callable[[RunConfig], Path] +EvaluateFn = Callable[[RunConfig, Path], dict[str, object]] +Log = Callable[[str], None] + +# Fields whose effect equivalent_run() compares through the rows and the +# step plan they produce instead of by value. +DATA_AND_SCHEDULE_FIELDS = frozenset( + {"per_intent", "k_shot", "oos_train", "num_train_epochs", "max_steps", "min_train_steps"} +) + + +def run_dir(config: RunConfig) -> Path: + return Path(config.checkpoint_root) / config.run_name + + +def archive_intact(paths: RunPaths, record: dict[str, object], seed: int) -> bool: + """The archive's SHA-256 agrees in results JSON, manifest and file, and it is this seed's.""" + if not paths.logits.exists(): + return False + logits = record.get("logits") + recorded = logits.get("sha256") if isinstance(logits, dict) else None + listed = read_manifest(paths.manifest).get(paths.logits.name, {}).get("sha256") + if not recorded or not recorded == listed == sha256_of(paths.logits): + return False + try: + archive = load_logits(paths.logits) + except ArchiveError: + return False + return archive.metadata["seed"] == seed + + +def read_record(paths: RunPaths) -> dict[str, object] | None: + if not paths.results_json.exists(): + return None + record = json.loads(paths.results_json.read_text(encoding="utf-8")) + return record if isinstance(record, dict) else None + + +def is_done(config: RunConfig) -> bool: + """This config's results JSON exists and its archive is intact and the one it was scored on.""" + paths = RunPaths.of(config) + record = read_record(paths) + if record is None or not config.matches(record.get("config")): + return False + return archive_intact(paths, record, config.seed) + + +def reusable_weights(config: RunConfig) -> bool: + """Final weights exist and were trained with exactly this config.""" + final = run_dir(config) / "final" + summary = final / "train_summary.json" + if not ((final / "model.safetensors").exists() and summary.exists()): + return False + return config.matches(json.loads(summary.read_text(encoding="utf-8")).get("config")) + + +def clear_outputs(config: RunConfig, log: Log) -> None: + """Remove this run's weights, results JSON, logits archive and manifest entry.""" + paths = RunPaths.of(config) + removed = [str(p) for p in (paths.results_json, paths.logits) if p.exists()] + for name in removed: + Path(name).unlink() + if remove_from_manifest(paths.manifest, paths.logits.name): + removed.append(f"manifest entry {paths.logits.name}") + if run_dir(config).exists(): + shutil.rmtree(run_dir(config)) + removed.append(str(run_dir(config))) + if removed: + log(f"{config.run_name}: cleared {', '.join(removed)}") + + +def run_one( + config: RunConfig, train_fn: TrainFn, evaluate_fn: EvaluateFn, log: Log, keep_weights: bool +) -> None: + """Train (or reuse weights) and evaluate unless done; delete weights unless kept.""" + if is_done(config): + log(f"{config.run_name}: results and logits for this config already archived, skipping") + else: + if reusable_weights(config): + log(f"{config.run_name}: reusing weights trained with this config") + final_dir = run_dir(config) / "final" + else: + clear_outputs(config, log) + log(f"{config.run_name}: training") + final_dir = train_fn(config) + evaluate_fn(config, final_dir) + log(f"{config.run_name}: evaluated, logits archived") + if not keep_weights and run_dir(config).exists(): + shutil.rmtree(run_dir(config)) + log(f"{config.run_name}: deleted weights {run_dir(config)}") + + +def config_from_record(recorded: object) -> RunConfig | None: + """Rebuild the RunConfig a results record was produced with; None if it cannot be exact.""" + if not isinstance(recorded, dict): + return None + filled = {**ADDED_FIELD_DEFAULTS, **recorded} + names = {f.name for f in fields(RunConfig)} + if set(filled) != names: + return None + try: + return RunConfig(**filled) + except (TypeError, ValueError): + return None + + +def data_key(config: RunConfig) -> tuple[int, int] | None: + """(k, oos rows) of the training sample; None for a per-intent cap (not compared this way). + + No sampling at all is the k=100 sample: CLINC150's train split has + exactly 100 rows per intent and 250 oos rows (pinned in SPLIT_FILES and + docs/DATA.md), and ``sample_k_shot`` at k=100 returns every row in the + original order (tests/test_sampling.py), so the Trainer sees the same + dataset. + """ + if config.per_intent is not None: + return None + if config.k_shot is None: + if full_train_rows() != planned_rows(FULL_K)[0]: + return None + return FULL_K, OOS_TRAIN_ROWS[FULL_K] + return config.k_shot, planned_rows(config.k_shot, config.oos_train)[1] + + +def training_identity(config: RunConfig) -> dict[str, object] | None: + """What the Trainer actually receives: rows, step count and every other field.""" + data = data_key(config) + if data is None: + return None + plan = plan_steps(config, planned_rows(*data)[0]) + if plan.max_steps_arg < 0: + schedule: tuple[str, float] = ("epochs", config.num_train_epochs) + else: + schedule = ("steps", plan.planned_steps) + rest = { + name: value + for name, value in config.identity().items() + if name not in DATA_AND_SCHEDULE_FIELDS | LOCATION_FIELDS + } + return {**rest, "data": data, "schedule": schedule} + + +def equivalent_run(config: RunConfig, record: dict[str, object]) -> bool: + """Whether ``record`` is a run that ``config`` would have reproduced step for step. + + Two rewrites are accepted, and nothing else: + + 1. no sampling (``k_shot`` and ``per_intent`` null) is the same data as + ``k_shot=100`` with the table's 250 oos rows (see ``data_key``); + 2. a ``min_train_steps`` at or below the epoch step count does not + change the run (``plan_steps`` keeps ``max_steps=-1`` then), and a + fixed step count reached through ``max_steps`` or through + ``min_train_steps`` is the same TrainingArguments. + + All other fields must be equal. The record's own training summary must + also show the rows and steps this config plans: train rows, oos rows, + and global_step. + """ + recorded = config_from_record(record.get("config")) + training = record.get("training") + if recorded is None or not isinstance(training, dict): + return False + mine, theirs = training_identity(config), training_identity(recorded) + if mine is None or theirs is None or mine != theirs: + return False + k, oos = mine["data"] # type: ignore[misc] + rows = planned_rows(k, oos)[0] + return ( + training.get("train_rows") == rows + and training.get("oos_train_rows") == oos + and training.get("global_step") == plan_steps(config, rows).planned_steps + and training.get("seed") == config.seed + ) + + +def default_train(config: RunConfig) -> Path: + from tinyrouter.train import prepare_train_split, train + + return train(config, prepare_train_split(config), run_dir(config)) + + +def default_evaluate(config: RunConfig, model_dir: Path) -> dict[str, object]: + from tinyrouter.evaluate import evaluate + + return evaluate(config, model_dir) + + +def reusable_equivalent(config: RunConfig, donor: RunConfig) -> dict[str, object] | None: + """The donor run's record if it is done, intact, and equivalent to ``config``; else None.""" + paths = RunPaths.of(donor) + record = read_record(paths) + if record is None or record.get("run_name") != donor.run_name: + return None + if not equivalent_run(config, record): + return None + return record if archive_intact(paths, record, config.seed) else None diff --git a/src/tinyrouter/sampling.py b/src/tinyrouter/sampling.py new file mode 100644 index 0000000..1b9aa7f --- /dev/null +++ b/src/tinyrouter/sampling.py @@ -0,0 +1,97 @@ +"""k-shot training samples for the learning curves (docs/PLAN.md section 4, AC5). + +Every in-scope intent gets exactly ``k`` rows. Out-of-scope gets the +number in ``OOS_TRAIN_ROWS``, a hardcoded table (not computed at run +time) equal to ceil(2.5k): the OOS+ train split's ratio of 250 oos rows +to 100 per intent, so ``k = 100`` is the whole train split. The table is +checked against ceil(2.5k) by a test, not by the code that uses it. + +The sample depends only on the split's content, ``k``, the oos count and +``seed``; the selected rows keep their original order. +""" + +from __future__ import annotations + +import hashlib + +import numpy as np + +from tinyrouter.data import SPLIT_FILES, Split, load_split +from tinyrouter.labels import load_label_space + +OOS_TRAIN_ROWS: dict[int, int] = {1: 3, 5: 13, 10: 25, 25: 63, 50: 125, 100: 250} +CURVE_KS: tuple[int, ...] = tuple(OOS_TRAIN_ROWS) +FULL_K = 100 +IN_SCOPE_INTENTS = 150 + + +class SamplingError(ValueError): + """A k-shot sample cannot be drawn as specified.""" + + +def oos_rows_for(k: int, oos_override: int | None = None) -> int: + """Out-of-scope rows in the k-shot sample; ``oos_override`` replaces the table (ablation).""" + if k not in OOS_TRAIN_ROWS: + raise SamplingError(f"k={k} is not a curve point; allowed: {CURVE_KS}") + if oos_override is None: + return OOS_TRAIN_ROWS[k] + if oos_override < 0: + raise SamplingError(f"oos override must be >= 0, got {oos_override}") + return oos_override + + +def planned_rows(k: int, oos_override: int | None = None) -> tuple[int, int]: + """(train rows, oos rows) a k-shot sample will have, without reading any data.""" + oos = oos_rows_for(k, oos_override) + return IN_SCOPE_INTENTS * k + oos, oos + + +def k_shot_indices( + intents: np.ndarray, k: int, seed: int, oos_override: int | None = None +) -> np.ndarray: + """Sorted row indices of the k-shot sample of a train split's intent column.""" + oos_count = oos_rows_for(k, oos_override) + oos_id = load_label_space().oos_intent_id + present = set(np.unique(intents).tolist()) + expected = set(range(load_label_space().num_intents)) + if present != expected: + raise SamplingError( + f"train split has {len(present)} distinct intents, expected all {len(expected)}" + ) + rng = np.random.default_rng(seed) + keep: list[np.ndarray] = [] + for intent in sorted(present): + rows = np.flatnonzero(intents == intent) + want = oos_count if intent == oos_id else k + if len(rows) < want: + raise SamplingError(f"intent {intent} has {len(rows)} rows, need {want}") + keep.append(rows if len(rows) == want else rng.choice(rows, size=want, replace=False)) + return np.sort(np.concatenate(keep)).astype(np.int64) + + +def sample_k_shot(split: Split, k: int, seed: int, oos_override: int | None = None) -> Split: + """The k-shot training sample; at k=100 with the table's 250 oos rows it is the whole split.""" + if split.name != "train": + raise SamplingError(f"only the train split is subsampled, got '{split.name}'") + return split.take(k_shot_indices(split.intents, k, seed, oos_override)) + + +def curve_sample(k: int, seed: int, oos_override: int | None = None) -> Split: + """The training rows of curve point (k, seed): shared by the encoders and the baselines.""" + return sample_k_shot(load_split("train"), k, seed, oos_override) + + +def sample_fingerprint(split: Split) -> str: + """SHA-256 over the sample's (text, intent) rows in order. + + Two runs with the same fingerprint trained on the same rows in the same + order, so an encoder and a baseline can be checked to share a sample. + """ + digest = hashlib.sha256() + for text, intent in zip(split.texts, split.intents.tolist(), strict=True): + digest.update(f"{intent}\t{text}\n".encode()) + return digest.hexdigest() + + +def full_train_rows() -> int: + return SPLIT_FILES["train"].rows diff --git a/src/tinyrouter/steps.py b/src/tinyrouter/steps.py new file mode 100644 index 0000000..661888c --- /dev/null +++ b/src/tinyrouter/steps.py @@ -0,0 +1,57 @@ +"""How many optimizer steps a run takes, decided before training starts. + +The learning-curve protocol (docs/PLAN.md section 4, hyperparameter +protocol) trains for max(S_min, the steps of ``num_train_epochs`` epochs). +Small k-shot samples give only a handful of steps per epoch, so S_min keeps +the tiny points from being judged on an undertrained model; the large +points are unaffected because their epoch steps already exceed S_min. + +The epoch count reproduces HF Trainer's own arithmetic (transformers +5.17, ``Trainer.set_initial_training_values``, one device, no gradient +accumulation): ceil(epochs * ceil(rows / batch)). ``train()`` checks the +step count it actually ran against this plan and raises on a mismatch. +""" + +from __future__ import annotations + +import math +from dataclasses import asdict, dataclass +from typing import Literal + +from tinyrouter.config import RunConfig + +DecidedBy = Literal["max_steps", "epochs", "min_train_steps"] + + +@dataclass(frozen=True) +class StepPlan: + epoch_steps: int | None + min_train_steps: int | None + planned_steps: int + decided_by: DecidedBy + # What TrainingArguments.max_steps receives: -1 means "train by epochs". + max_steps_arg: int + + def as_dict(self) -> dict[str, object]: + return asdict(self) + + +def epoch_steps(train_rows: int, batch_size: int, epochs: float) -> int: + if train_rows < 1: + raise ValueError("cannot plan steps for an empty training set") + return math.ceil(epochs * math.ceil(train_rows / batch_size)) + + +def plan_steps(config: RunConfig, train_rows: int) -> StepPlan: + """Steps for ``config`` on ``train_rows`` rows; ties between the two rules go to epochs. + + Going to epochs on a tie keeps the code path of a run without + ``min_train_steps`` (``max_steps=-1``), so such a run is the same run. + """ + if config.max_steps > 0: + return StepPlan(None, None, config.max_steps, "max_steps", config.max_steps) + by_epochs = epoch_steps(train_rows, config.train_batch_size, config.num_train_epochs) + floor = config.min_train_steps + if floor is not None and floor > by_epochs: + return StepPlan(by_epochs, floor, floor, "min_train_steps", floor) + return StepPlan(by_epochs, floor, by_epochs, "epochs", -1) diff --git a/src/tinyrouter/train.py b/src/tinyrouter/train.py index bd71875..17f01ac 100644 --- a/src/tinyrouter/train.py +++ b/src/tinyrouter/train.py @@ -24,6 +24,8 @@ start_measurement, ) from tinyrouter.labels import LabelSpace, load_label_space +from tinyrouter.sampling import curve_sample, sample_fingerprint +from tinyrouter.steps import StepPlan, plan_steps SUMMARY_NAME = "train_summary.json" @@ -70,9 +72,26 @@ def load_model_and_tokenizer(config: RunConfig, labels: LabelSpace) -> tuple[obj return model, tokenizer -def training_arguments(config: RunConfig, output_dir: Path, device: str) -> object: +def require_learning_rate(config: RunConfig) -> None: + if config.learning_rate is None: + raise ValueError( + f"{config.model_name}: learning_rate is null (not chosen yet); run `make pilot-lr` " + "and write the selected value into configs/curve.yaml" + ) + + +class StepCountError(RuntimeError): + """Training ran a different number of steps than the plan recorded for it.""" + + +def training_arguments( + config: RunConfig, output_dir: Path, device: str, plan: StepPlan | None = None +) -> object: + """TrainingArguments for ``config``; ``plan`` supplies max_steps (default: config.max_steps).""" from transformers import TrainingArguments + require_learning_rate(config) + max_steps = config.max_steps if plan is None else plan.max_steps_arg return TrainingArguments( output_dir=str(output_dir), seed=config.seed, @@ -82,14 +101,14 @@ def training_arguments(config: RunConfig, output_dir: Path, device: str) -> obje # transformers 5.x reads warmup_steps < 1 as a fraction of total steps. warmup_steps=config.warmup_ratio, num_train_epochs=config.num_train_epochs, - max_steps=config.max_steps, + max_steps=max_steps, per_device_train_batch_size=config.train_batch_size, per_device_eval_batch_size=config.eval_batch_size, eval_strategy="no", - save_strategy="epoch" if config.max_steps < 0 else "no", + save_strategy="epoch" if max_steps < 0 else "no", save_total_limit=1, save_only_model=True, - logging_strategy="epoch" if config.max_steps < 0 else "steps", + logging_strategy="epoch" if max_steps < 0 else "steps", logging_steps=1, report_to="none", use_cpu=device == "cpu", @@ -102,13 +121,16 @@ def train(config: RunConfig, train_split: Split, output_dir: Path) -> Path: """Train on ``train_split`` and return the directory holding the final weights. ``final/train_summary.json`` records the training set size, parameter - counts, wall time, steps and peak memory; evaluation copies it into the - results JSON, because the weights (and this file) may be deleted later. + counts, wall time, the step plan and the steps actually run, and peak + memory; evaluation copies it into the results JSON, because the weights + (and this file) may be deleted later. """ from transformers import DataCollatorWithPadding, Trainer, set_seed if train_split.name != "train": raise ValueError(f"train() only accepts the train split, got '{train_split.name}'") + require_learning_rate(config) + plan = plan_steps(config, len(train_split)) set_seed(config.seed) labels = load_label_space() device = pick_device(config.device) @@ -116,7 +138,7 @@ def train(config: RunConfig, train_split: Split, output_dir: Path) -> Path: sampler = start_measurement(device) trainer = Trainer( model=model, # type: ignore[arg-type] - args=training_arguments(config, output_dir, device), # type: ignore[arg-type] + args=training_arguments(config, output_dir, device, plan), # type: ignore[arg-type] train_dataset=to_hf_dataset(train_split, tokenizer, config.max_length), # type: ignore[arg-type] data_collator=DataCollatorWithPadding(tokenizer), # type: ignore[arg-type] processing_class=tokenizer, # type: ignore[arg-type] @@ -125,6 +147,12 @@ def train(config: RunConfig, train_split: Split, output_dir: Path) -> Path: started = time.perf_counter() result = trainer.train() wall_seconds = time.perf_counter() - started + if result.global_step != plan.planned_steps: + raise StepCountError( + f"{config.run_name}: trained {result.global_step} steps, planned " + f"{plan.planned_steps} ({plan.decided_by}); steps.plan_steps no longer matches the " + "Trainer's arithmetic" + ) sampler.sample() final_dir = output_dir / "final" trainer.save_model(str(final_dir)) @@ -141,6 +169,9 @@ def train(config: RunConfig, train_split: Split, output_dir: Path) -> Path: "train_rows": len(train_split), "oos_train_rows": int((train_split.intents == labels.oos_intent_id).sum()), "per_intent": config.per_intent, + "k_shot": config.k_shot, + "train_sample_sha256": sample_fingerprint(train_split), + "step_plan": plan.as_dict(), "global_step": result.global_step, "train_loss": result.training_loss, "parameters": count_parameters(model), @@ -152,6 +183,9 @@ def train(config: RunConfig, train_split: Split, output_dir: Path) -> Path: def prepare_train_split(config: RunConfig) -> Split: + """The rows ``config`` trains on: a k-shot sample, a per-intent cap (smoke), or everything.""" + if config.k_shot is not None: + return curve_sample(config.k_shot, config.seed, config.oos_train) return subsample_per_intent(load_split("train"), config.per_intent, config.seed) diff --git a/tests/archive_fakes.py b/tests/archive_fakes.py index fbe48c6..481e996 100644 --- a/tests/archive_fakes.py +++ b/tests/archive_fakes.py @@ -2,6 +2,7 @@ import numpy as np +from tinyrouter.archive import pinned_dataset_sha256 from tinyrouter.calibrate import SplitLogits from tinyrouter.labels import load_label_space @@ -22,16 +23,18 @@ def fake_splits(seed: int = 0, n_val: int = 7, n_test: int = 11) -> dict[str, Sp def fake_metadata(**overrides: object) -> dict[str, object]: meta: dict[str, object] = { - "format_version": 1, + "format_version": 2, "run_name": "bert-base-uncased-full-seed42", "model_name": "google-bert/bert-base-uncased", "model_revision": "86b5e0934494bd15c9632b12f734a8a67f723594", "seed": 42, "per_intent": None, + "k_shot": None, "train_rows": 15_250, "oos_train_rows": 250, "eval_per_intent": None, "dataset_revision": "155b9c710419136e17307b80d0a13e68cd46b4ec", + "dataset_sha256": pinned_dataset_sha256(), "label_space_sha256": load_label_space().sha256, "git_commit": "0" * 40, "git_dirty": False, diff --git a/tests/run_fakes.py b/tests/run_fakes.py new file mode 100644 index 0000000..f370891 --- /dev/null +++ b/tests/run_fakes.py @@ -0,0 +1,112 @@ +"""A stand-in for train() and evaluate() that writes the same files, for pilot and curve tests. + +The training summary carries what the real one does for these tests: the +row counts ``planned_rows`` gives, the step plan and a ``global_step`` equal +to it. Evaluation writes a real archive, manifest entry and results JSON. +""" + +import json +from dataclasses import asdict +from pathlib import Path + +from archive_fakes import fake_metadata +from tinyrouter.archive import record_in_manifest, save_logits +from tinyrouter.calibrate import SplitLogits +from tinyrouter.config import RunConfig +from tinyrouter.evaluate import RunPaths, score +from tinyrouter.sampling import planned_rows +from tinyrouter.steps import plan_steps + + +def summary_for(config: RunConfig) -> dict[str, object]: + if config.k_shot is None: + rows, oos = 15_250, 250 + else: + rows, oos = planned_rows(config.k_shot, config.oos_train) + plan = plan_steps(config, rows) + return { + "run_name": config.run_name, + "seed": config.seed, + "config": asdict(config), + "train_rows": rows, + "oos_train_rows": oos, + "k_shot": config.k_shot, + "train_sample_sha256": f"sample-{config.k_shot}-{config.seed}", + "step_plan": plan.as_dict(), + "global_step": plan.planned_steps, + "train_wall_seconds": 1.0, + } + + +class FakeRuns: + def __init__(self) -> None: + self.trained: list[str] = [] + self.evaluated: list[str] = [] + + def train(self, config: RunConfig) -> Path: + self.trained.append(f"{config.run_name}@{config.checkpoint_root}") + final = Path(config.checkpoint_root) / config.run_name / "final" + final.mkdir(parents=True) + (final / "model.safetensors").write_bytes(b"weights") + (final / "train_summary.json").write_text(json.dumps(summary_for(config))) + return final + + def evaluate(self, config: RunConfig, model_dir: Path) -> dict[str, object]: + self.evaluated.append(config.run_name) + paths = RunPaths.of(config) + paths.results_json.unlink(missing_ok=True) + training = json.loads((model_dir / "train_summary.json").read_text()) + shifted = fake_validation(100 + config.seed + (config.k_shot or 0)) + test = SplitLogits("test", shifted.logits, shifted.labels) + splits = {"validation": self.validation(config, model_dir), "test": test} + meta = fake_metadata( + run_name=config.run_name, + seed=config.seed, + k_shot=config.k_shot, + train_rows=training["train_rows"], + oos_train_rows=training["oos_train_rows"], + ) + save_logits(paths.logits, splits, meta) + entry = record_in_manifest(paths.manifest, paths.logits) + record = { + "run_name": config.run_name, + "config": asdict(config), + "training": training, + "logits": {"file": paths.logits.name, "sha256": entry["sha256"]}, + "metrics": score(splits["validation"], splits["test"]), + } + paths.results_json.parent.mkdir(parents=True, exist_ok=True) + paths.results_json.write_text(json.dumps(record)) + return record + + def validation(self, config: RunConfig, model_dir: Path) -> SplitLogits: + """Validation logits whose accuracy the test controls through ``self.correct``. + + The seed shifts the count by one so different seeds give different archives. + """ + correct = getattr(self, "correct", {}).get(self.key(config), 0) + return fake_validation(correct + config.seed % 42) + + @staticmethod + def key(config: RunConfig) -> tuple[str, object]: + value = config.min_train_steps if config.k_shot == 5 else config.learning_rate + return config.model_name, value + + +def fake_validation(in_scope_correct: int, oos_correct: int = 0) -> SplitLogits: + """3,000 in-scope rows (20 per intent) and 100 oos rows, with chosen counts right.""" + import numpy as np + + from tinyrouter.labels import load_label_space + + oos = load_label_space().oos_intent_id + in_scope_ids = [i for i in range(151) if i != oos] + gold = np.array([i for i in in_scope_ids for _ in range(20)] + [oos] * 100, dtype=np.int64) + pred = np.where(gold == oos, (oos + 1) % 151, (gold + 1) % 151) + pred[np.flatnonzero(pred == oos)] = (oos + 2) % 151 + pred[:in_scope_correct] = gold[:in_scope_correct] + oos_rows = np.flatnonzero(gold == oos)[:oos_correct] + pred[oos_rows] = oos + logits = np.full((len(gold), 151), -3.0, dtype=np.float32) + logits[np.arange(len(gold)), pred] = 3.0 + return SplitLogits("validation", logits, gold) diff --git a/tests/test_archive.py b/tests/test_archive.py index b31b66b..ff4afd7 100644 --- a/tests/test_archive.py +++ b/tests/test_archive.py @@ -1,4 +1,5 @@ import json +from pathlib import Path import numpy as np import pytest @@ -193,3 +194,108 @@ def git(*args): assert git_state(tmp_path)[1] is False (tmp_path / "code.py").write_text("x = 2\n") assert git_state(tmp_path)[1] is True + + +def write_format1(path, **overrides): + """An archive as the AC2 code wrote it: format 1, no dataset_sha256, no k_shot.""" + meta = fake_metadata(format_version=1, **overrides) + del meta["dataset_sha256"], meta["k_shot"] + arrays = {"metadata": np.array(json.dumps(meta, sort_keys=True))} + for name, split in fake_splits().items(): + arrays[f"{name}_logits"] = split.logits + arrays[f"{name}_labels"] = split.labels.astype(np.int64) + with path.open("wb") as fh: + np.savez_compressed(fh, **arrays) + return path + + +def test_new_archives_record_the_sha256_of_all_three_split_files(tmp_path): + from tinyrouter.data import SPLIT_FILES + + archive = load_logits(save_logits(tmp_path / "r.npz", fake_splits(), fake_metadata())) + assert archive.metadata["format_version"] == 2 + assert archive.metadata["dataset_sha256"] == { + name: SPLIT_FILES[name].sha256 for name in ("train", "validation", "test") + } + assert "dataset_sha256_source" not in archive.metadata + + +def test_format1_archive_at_the_pinned_revision_is_back_filled_on_read(tmp_path): + from tinyrouter.archive import BACKFILL_NOTE, pinned_dataset_sha256 + + archive = load_logits(write_format1(tmp_path / "old.npz")) + assert archive.metadata["format_version"] == 1 + assert archive.metadata["dataset_sha256"] == pinned_dataset_sha256() + assert archive.metadata["dataset_sha256_source"] == BACKFILL_NOTE + assert archive.metadata["k_shot"] is None + + +def test_format1_archive_from_another_dataset_revision_is_refused(tmp_path): + path = write_format1(tmp_path / "old.npz", dataset_revision="f" * 40) + with pytest.raises(ArchiveError, match="can only be back-filled"): + load_logits(path) + + +def test_format1_archive_that_already_claims_checksums_is_refused(tmp_path): + path = write_format1(tmp_path / "old.npz") + rewrite_metadata(path, lambda m: m.update(dataset_sha256={"train": "x"})) + with pytest.raises(ArchiveError, match="already carries"): + load_logits(path) + + +@pytest.mark.parametrize( + "bad", + [ + {"train": "0" * 64, "validation": "0" * 64, "test": "0" * 64}, + {"train": "x", "validation": "y"}, + None, + ], +) +def test_format2_archive_with_other_dataset_checksums_is_refused(tmp_path, bad): + with pytest.raises(ArchiveError, match="dataset_sha256|pinned CLINC150"): + save_logits(tmp_path / "r.npz", fake_splits(), fake_metadata(dataset_sha256=bad)) + + +def test_format2_archive_of_another_revision_is_refused_even_with_pinned_checksums(tmp_path): + with pytest.raises(ArchiveError, match="pinned CLINC150"): + save_logits(tmp_path / "r.npz", fake_splits(), fake_metadata(dataset_revision="f" * 40)) + + +def test_format1_cannot_be_written_any_more(tmp_path): + with pytest.raises(ArchiveError, match="cannot be written"): + save_logits(tmp_path / "r.npz", fake_splits(), fake_metadata(format_version=1)) + + +def test_load_validation_logits_never_needs_the_test_arrays(tmp_path): + from tinyrouter.archive import load_validation_logits + + path = save_logits(tmp_path / "r.npz", fake_splits(), fake_metadata()) + with np.load(path, allow_pickle=False) as npz: + kept = {k: npz[k] for k in npz.files if not k.startswith("test_")} + with path.open("wb") as fh: + np.savez_compressed(fh, **kept) + with pytest.raises(ArchiveError, match="test_"): + load_logits(path) + meta, val = load_validation_logits(path) + assert val.split == "validation" and val.logits.shape == (7, NUM_INTENTS) + assert meta["run_name"] == "bert-base-uncased-full-seed42" + + +def test_load_validation_logits_back_fills_a_format1_archive(tmp_path): + from tinyrouter.archive import load_validation_logits + + meta, _ = load_validation_logits(write_format1(tmp_path / "old.npz")) + assert meta["dataset_sha256_source"] + + +AC2_ARCHIVES = sorted( + (Path(__file__).parent.parent / "results" / "logits").glob("bert-base-uncased-full-*.npz") +) + + +@pytest.mark.skipif(not AC2_ARCHIVES, reason="AC2 archives are not in git (GitHub Release)") +@pytest.mark.parametrize("path", AC2_ARCHIVES, ids=lambda p: p.name) +def test_the_real_ac2_archives_still_load_as_format1(path): + archive = load_logits(path) + assert archive.metadata["format_version"] == 1 + assert archive.test.logits.shape == (5_500, NUM_INTENTS) diff --git a/tests/test_baselines.py b/tests/test_baselines.py new file mode 100644 index 0000000..8034984 --- /dev/null +++ b/tests/test_baselines.py @@ -0,0 +1,129 @@ +import json +from collections import Counter + +import numpy as np +import pytest + +from test_sampling import FULL +from tinyrouter import baselines +from tinyrouter.archive import check_against_manifest, load_logits +from tinyrouter.config import RunConfig +from tinyrouter.data import Split +from tinyrouter.evaluate import RunPaths, score +from tinyrouter.labels import AGENTS, load_label_space +from tinyrouter.sampling import OOS_TRAIN_ROWS, sample_fingerprint, sample_k_shot + +pytestmark = pytest.mark.filterwarnings("ignore::tinyrouter.calibrate.TemperatureBoundWarning") + +LABELS = load_label_space() +OOS = LABELS.oos_intent_id + + +def worded(split: Split) -> Split: + """Give each intent its own vocabulary so a text is closest to its own intent's centroid.""" + texts = tuple(f"w{i}a w{i}b common n{n % 7}" for n, i in enumerate(split.intents)) + return Split(split.name, texts, split.intents) + + +TRAIN = worded(FULL) + + +def eval_split(name: str) -> Split: + """Validation or test: every intent twice, oos five times.""" + intents = np.array([i for i in range(151) for _ in range(5 if i == OOS else 2)]) + return worded(Split(name, tuple("" for _ in intents), intents)) # type: ignore[arg-type] + + +EVALS = {name: eval_split(name) for name in ("validation", "test")} + + +def test_majority_scores_are_the_same_prior_for_every_row(): + train = sample_k_shot(TRAIN, 5, seed=42) + scores = baselines.majority_scores(train, 4) + assert scores.shape == (4, 151) and scores.dtype == np.float32 + assert (scores == scores[0]).all() + assert np.exp(scores[0].astype(np.float64)).sum() == pytest.approx(1.0, rel=1e-6) + assert np.exp(scores[0, OOS]) == pytest.approx(OOS_TRAIN_ROWS[5] / len(train), rel=1e-6) + + +def test_majority_argmax_is_oos_and_its_agent_sum_is_finance(): + scores = baselines.majority_scores(sample_k_shot(TRAIN, 10, seed=42), 1) + assert scores.argmax() == OOS + probs = np.exp(scores.astype(np.float64)) + assert AGENTS[LABELS.aggregate_probs(probs).argmax()] == "finance_agent" + + +def test_majority_refuses_a_sample_missing_an_intent(): + train = sample_k_shot(TRAIN, 5, seed=42) + keep = np.flatnonzero(train.intents != 9) + with pytest.raises(ValueError, match=r"\[9\]"): + baselines.majority_scores(train.take(keep), 1) + + +def test_tfidf_centroid_puts_each_query_nearest_its_own_intent(): + model = baselines.TfidfCentroid(sample_k_shot(TRAIN, 5, seed=42)) + scores = model.scores(EVALS["validation"].texts) + assert (scores.argmax(axis=1) == EVALS["validation"].intents).all() + assert scores.max() <= baselines.COSINE_SCALE + 1e-3 and scores.min() >= 0 + + +def test_tfidf_is_fit_on_the_training_sample_only(): + model = baselines.TfidfCentroid(sample_k_shot(TRAIN, 1, seed=42)) + unseen = model.scores(("zzz_never_seen_in_training",)) + assert (unseen == 0).all() + + +def test_baseline_run_is_archived_and_scored_like_an_encoder_run(tmp_path): + train = sample_k_shot(TRAIN, 5, seed=43) + record = baselines.run_baseline("tfidf-centroid", 5, 43, train, EVALS, tmp_path) + paths = RunPaths.named(tmp_path, "tfidf-centroid-k5-seed43") + check_against_manifest(paths.manifest, paths.logits) + archive = load_logits(paths.logits) + meta = archive.metadata + assert (meta["model_name"], meta["k_shot"], meta["seed"]) == ("baseline/tfidf-centroid", 5, 43) + assert meta["train_rows"] == 150 * 5 + 13 and meta["oos_train_rows"] == 13 + assert archive.validation.logits.shape == (len(EVALS["validation"]), 151) + expected = score(archive.validation, archive.test) + assert set(record["metrics"]) == set(expected) + assert set(record["metrics"]["test"]["raw"]) == set(expected["test"]["raw"]) + assert ( + json.loads(paths.results_json.read_text())["logits"]["sha256"] == record["logits"]["sha256"] + ) + + +def test_baseline_rerun_on_the_same_sample_does_nothing(tmp_path, monkeypatch): + train = sample_k_shot(TRAIN, 5, seed=42) + baselines.run_baseline("majority", 5, 42, train, EVALS, tmp_path) + + def refit(*_): + raise AssertionError("an intact run with the same sample must not be recomputed") + + monkeypatch.setattr(baselines, "baseline_logits", refit) + baselines.run_baseline("majority", 5, 42, train, EVALS, tmp_path) + other = sample_k_shot(TRAIN, 5, seed=44) + with pytest.raises(AssertionError, match="must not be recomputed"): + baselines.run_baseline("majority", 5, 42, other, EVALS, tmp_path) + + +def test_baselines_see_exactly_the_rows_the_encoder_of_that_point_trains_on(tmp_path, monkeypatch): + import tinyrouter.sampling as sampling + from tinyrouter.train import prepare_train_split + + monkeypatch.setattr(sampling, "load_split", lambda name: TRAIN) + body = baselines.run_all(tmp_path, eval_fn=EVALS.__getitem__, log=lambda _: None) + assert len(body["points"]) == 6 * 3 * 2 + for point in body["points"]: + encoder = RunConfig( + model_name="m", model_revision="r", k_shot=point["k"], seed=point["seed"] + ) + assert point["train_sample_sha256"] == sample_fingerprint(prepare_train_split(encoder)) + counts = Counter((p["baseline"], p["k"]) for p in body["points"]) + assert set(counts.values()) == {3} + shas = [p["logits_sha256"] for p in body["points"]] + assert len(set(shas)) == len(shas) + assert (tmp_path / "curves" / "baselines.json").exists() + + +def test_unknown_baseline_is_refused(): + with pytest.raises(ValueError, match="unknown baseline"): + baselines.baseline_logits("knn", TRAIN, EVALS) diff --git a/tests/test_config.py b/tests/test_config.py index b3e3de1..6c6290e 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -1,10 +1,14 @@ +from dataclasses import asdict, replace from pathlib import Path import pytest from tinyrouter.config import load_config -CONFIGS = sorted((Path(__file__).parent.parent / "configs").glob("*.yaml")) +# curve.yaml is the curve protocol (tests/test_protocol.py), not a run config. +CONFIGS = sorted( + p for p in (Path(__file__).parent.parent / "configs").glob("*.yaml") if p.name != "curve.yaml" +) @pytest.mark.parametrize("path", CONFIGS, ids=lambda p: p.name) @@ -30,5 +34,23 @@ def test_warmup_ratio_of_one_or_more_is_rejected(tmp_path): def test_run_name_encodes_model_size_and_seed(tmp_path): path = tmp_path / "c.yaml" path.write_text("model_name: org/bert-x\nmodel_revision: r\nper_intent: 16\nseed: 43\n") - assert load_config(path).run_name == "bert-x-k16-seed43" - assert load_config(path).with_seed(44).run_name == "bert-x-k16-seed44" + assert load_config(path).run_name == "bert-x-cap16-seed43" + assert load_config(path).with_seed(44).run_name == "bert-x-cap16-seed44" + + +def test_run_name_encodes_k_shot_and_the_oos_override(): + from tinyrouter.config import RunConfig + + config = RunConfig(model_name="answerdotai/ModernBERT-base", model_revision="r", k_shot=100) + assert config.run_name == "ModernBERT-base-k100-seed42" + assert replace(config, oos_train=0).run_name == "ModernBERT-base-k100-oos0-seed42" + + +def test_a_record_written_before_the_step3_fields_matches_their_defaults(): + from tinyrouter.config import ADDED_FIELD_DEFAULTS, RunConfig + + config = RunConfig(model_name="m", model_revision="r") + old = {k: v for k, v in asdict(config).items() if k not in ADDED_FIELD_DEFAULTS} + assert config.matches(old) + assert not replace(config, min_train_steps=100).matches(old) + assert not replace(config, k_shot=5).matches(old) diff --git a/tests/test_curves.py b/tests/test_curves.py new file mode 100644 index 0000000..45f28b9 --- /dev/null +++ b/tests/test_curves.py @@ -0,0 +1,194 @@ +import json +from dataclasses import replace +from pathlib import Path + +import pytest + +from run_fakes import FakeRuns +from test_pilots import make_protocol +from tinyrouter.curves import CurveError, index_entry, run_curve +from tinyrouter.evaluate import RunPaths +from tinyrouter.protocol import ProtocolError +from tinyrouter.runs import run_one + +pytestmark = pytest.mark.filterwarnings("ignore::tinyrouter.calibrate.TemperatureBoundWarning") + + +def quiet(_: str) -> None: + pass + + +def ac2_runs(protocol) -> None: + """The three AC2 runs, made with BERT's own config file (lr 5e-5, no k_shot, no S_min).""" + fake = FakeRuns() + for seed in (42, 43, 44): + config = protocol.base("bert").with_seed(seed) + run_one(config, fake.train, fake.evaluate, quiet, keep_weights=seed == 42) + + +def by_point(body) -> dict: + return {(p["k"], p["seed"]): p for p in body["points"]} + + +def test_bert_curve_reuses_ac2_at_k100_and_trains_the_other_fifteen(tmp_path): + protocol = make_protocol(tmp_path, s_min=400, bert_lr=5e-5, modern_lr=2e-5) + ac2_runs(protocol) + fake = FakeRuns() + body = run_curve( + "bert", protocol.curve_configs("bert"), protocol, fake.train, fake.evaluate, quiet + ) + points = by_point(body) + assert len(points) == 18 and len(fake.trained) == 15 + assert not any("k100" in name for name in fake.trained) + for seed in (42, 43, 44): + full = points[(100, seed)] + assert full["reused_from"] == f"bert-base-uncased-full-seed{seed}" + assert full["run_name"] == full["reused_from"] + assert (full["planned_steps"], full["decided_by"], full["global_step"]) == ( + 2385, + "epochs", + 2385, + ) + assert points[(1, 42)]["decided_by"] == "min_train_steps" + assert points[(1, 42)]["planned_steps"] == points[(1, 42)]["global_step"] == 400 + assert points[(1, 42)]["epoch_steps"] == 25 + assert points[(50, 43)]["decided_by"] == "epochs" # 7,625 rows -> 1,195 steps + on_disk = json.loads((tmp_path / "res" / "curves" / "bert.json").read_text()) + assert on_disk == json.loads(json.dumps(body)) + + +def test_bert_curve_trains_k100_itself_when_the_pilot_picked_another_lr(tmp_path): + protocol = make_protocol(tmp_path, s_min=400, bert_lr=2e-5, modern_lr=2e-5) + ac2_runs(protocol) + fake = FakeRuns() + body = run_curve( + "bert", protocol.curve_configs("bert"), protocol, fake.train, fake.evaluate, quiet + ) + assert len(fake.trained) == 18 + assert all(p["reused_from"] is None for p in body["points"]) + assert by_point(body)[(100, 42)]["run_name"] == "bert-base-uncased-k100-seed42" + + +def test_an_ac2_seed_whose_archive_changed_is_retrained_not_reused(tmp_path): + protocol = make_protocol(tmp_path, s_min=400, bert_lr=5e-5, modern_lr=2e-5) + ac2_runs(protocol) + archive = RunPaths.of(protocol.base("bert").with_seed(43)).logits + archive.write_bytes(archive.read_bytes() + b"x") + configs = [c for c in protocol.curve_configs("bert") if c.k_shot == 100] + fake = FakeRuns() + body = run_curve("bert", configs, protocol, fake.train, fake.evaluate, quiet) + assert [p["reused_from"] is not None for p in body["points"]] == [True, False, True] + + +def test_curve_keeps_no_weights_and_leaves_ac2_seed_42_weights(tmp_path): + protocol = make_protocol(tmp_path, s_min=200, bert_lr=5e-5, modern_lr=2e-5) + ac2_runs(protocol) + fake = FakeRuns() + run_curve("bert", protocol.curve_configs("bert"), protocol, fake.train, fake.evaluate, quiet) + kept = sorted(p.parent.parent.name for p in (tmp_path / "ckpt").rglob("model.safetensors")) + assert kept == ["bert-base-uncased-full-seed42"] + + +def test_rerunning_a_finished_curve_trains_nothing(tmp_path): + protocol = make_protocol(tmp_path, s_min=200, bert_lr=5e-5, modern_lr=2e-5) + configs = protocol.curve_configs("modernbert") + run_curve("modernbert", configs, protocol, FakeRuns().train, FakeRuns().evaluate, quiet) + again = FakeRuns() + run_curve("modernbert", configs, protocol, again.train, again.evaluate, quiet) + assert again.trained == [] + + +def test_changing_s_min_redoes_every_point_instead_of_mixing_protocols(tmp_path): + protocol = make_protocol(tmp_path, s_min=200, bert_lr=5e-5, modern_lr=2e-5) + run_curve( + "modernbert", + protocol.curve_configs("modernbert"), + protocol, + FakeRuns().train, + FakeRuns().evaluate, + quiet, + ) + moved = replace(protocol, min_train_steps=400) + again = FakeRuns() + run_curve( + "modernbert", moved.curve_configs("modernbert"), moved, again.train, again.evaluate, quiet + ) + # S_min is part of every run's identity, so every point is redone; the index + # never mixes points made under two protocols. + assert len(again.trained) == 18 + + +def test_ablation_runs_modernbert_k100_without_oos_rows_and_never_reuses_ac2(tmp_path): + protocol = make_protocol(tmp_path, s_min=400, bert_lr=5e-5, modern_lr=2e-5) + ac2_runs(protocol) + fake = FakeRuns() + body = run_curve( + "oos-ablation", protocol.ablation_configs(), protocol, fake.train, fake.evaluate, quiet + ) + assert len(fake.trained) == 3 + assert {p["oos_train_rows"] for p in body["points"]} == {0} + assert {p["train_rows"] for p in body["points"]} == {15_000} + assert body["points"][0]["run_name"] == "ModernBERT-base-k100-oos0-seed42" + assert (tmp_path / "res" / "curves" / "oos-ablation.json").exists() + + +def test_curve_refuses_to_start_before_the_pilots_are_copied_in(tmp_path): + protocol = make_protocol(tmp_path, bert_lr=5e-5) + with pytest.raises(ProtocolError): + protocol.curve_configs("bert") + + +def test_a_point_that_ran_another_step_count_is_refused(tmp_path): + protocol = make_protocol(tmp_path, s_min=400, bert_lr=5e-5, modern_lr=2e-5) + config = protocol.curve_config("bert", 5, 42) + record = { + "run_name": config.run_name, + "logits": {"file": "x.npz", "sha256": "0" * 64}, + "training": {"train_rows": 763, "oos_train_rows": 13, "global_step": 120}, + } + with pytest.raises(CurveError, match="ran 120 steps"): + index_entry(config, record, None) + + +def cli_fakes(monkeypatch, protocol) -> list[str]: + import tinyrouter.curves as curves + + calls: list[str] = [] + fake = FakeRuns() + monkeypatch.setattr(curves, "load_protocol", lambda _: protocol) + monkeypatch.setattr(curves, "run_baselines", lambda root: calls.append(f"baselines {root}")) + + def curve(name, configs, proto): + calls.append(f"curve {name}") + return run_curve(name, configs, proto, fake.train, fake.evaluate, quiet) + + monkeypatch.setattr(curves, "run_curve", curve) + return calls + + +def test_the_cli_runs_baselines_before_the_curve(tmp_path, monkeypatch): + import tinyrouter.curves as curves + + protocol = make_protocol(tmp_path, s_min=400, bert_lr=5e-5, modern_lr=2e-5) + calls = cli_fakes(monkeypatch, protocol) + curves.main(["--model", "modernbert"]) + assert calls == [f"baselines {tmp_path / 'res'}", "curve modernbert"] + assert (tmp_path / "res" / "curves" / "modernbert.json").exists() + + +def test_the_cli_stops_before_the_baselines_while_a_pilot_value_is_missing(tmp_path, monkeypatch): + import tinyrouter.curves as curves + + calls = cli_fakes(monkeypatch, make_protocol(tmp_path, bert_lr=5e-5, modern_lr=2e-5)) + with pytest.raises(ProtocolError, match="pilot-steps"): + curves.main(["--model", "bert"]) + assert calls == [] + + +def test_the_ablation_cli_skips_the_baselines(tmp_path, monkeypatch): + import tinyrouter.curves as curves + + calls = cli_fakes(monkeypatch, make_protocol(tmp_path, s_min=400, bert_lr=5e-5, modern_lr=2e-5)) + curves.main(["--ablation"]) + assert calls == ["curve oos-ablation"] + assert Path(tmp_path / "res" / "curves" / "oos-ablation.json").exists() diff --git a/tests/test_metrics.py b/tests/test_metrics.py index 5dd6bea..2bdaa8b 100644 --- a/tests/test_metrics.py +++ b/tests/test_metrics.py @@ -11,6 +11,7 @@ oos_recall, routing_metrics, softmax, + wilson_interval, ) OOS_ID = 9 @@ -109,3 +110,42 @@ def test_routing_metrics_on_a_tiny_label_space(): assert m["oos_recall_8_summed"] == 0.0 assert m["n"] == 4 assert len(AGENTS) == 8 + + +@pytest.mark.parametrize( + ("successes", "n", "low", "high"), + [ + (0, 10, 0.0, 0.2775), # upper = z^2 / (n + z^2) + (5, 10, 0.2366, 0.7634), + (10, 10, 0.7225, 1.0), + ], +) +def test_wilson_interval_matches_hand_computed_values(successes, n, low, high): + got = wilson_interval(successes, n) + assert got == pytest.approx((low, high), abs=1e-4) + + +@pytest.mark.parametrize(("successes", "n"), [(583, 1000), (1, 3), (59, 100), (990, 1000)]) +def test_wilson_interval_agrees_with_scipy(successes, n): + from scipy.stats import binomtest, norm + + ref = binomtest(successes, n).proportion_ci(confidence_level=0.95, method="wilson") + got = wilson_interval(successes, n, z=float(norm.ppf(0.975))) + assert got == pytest.approx((ref.low, ref.high), abs=1e-12) + + +def test_wilson_interval_is_mirror_symmetric_and_narrows_with_n(): + low, high = wilson_interval(30, 100) + mirror = wilson_interval(70, 100) + assert (1 - high, 1 - low) == pytest.approx(mirror) + small, big = wilson_interval(58, 100), wilson_interval(583, 1000) + assert big[1] - big[0] < small[1] - small[0] + assert wilson_interval(583, 1000, z=2.576)[0] < big[0] + + +@pytest.mark.parametrize( + ("successes", "n", "z"), [(1, 0, 1.96), (-1, 10, 1.96), (11, 10, 1.96), (5, 10, 0.0)] +) +def test_wilson_interval_refuses_impossible_inputs(successes, n, z): + with pytest.raises(ValueError): + wilson_interval(successes, n, z) diff --git a/tests/test_pilots.py b/tests/test_pilots.py new file mode 100644 index 0000000..dc4ca08 --- /dev/null +++ b/tests/test_pilots.py @@ -0,0 +1,218 @@ +import json +from dataclasses import replace +from pathlib import Path + +import numpy as np +import pytest + +from archive_fakes import fake_splits +from run_fakes import FakeRuns, fake_validation +from tinyrouter import pilots +from tinyrouter.calibrate import LeakageError, SplitLogits +from tinyrouter.config import load_config +from tinyrouter.data import Split +from tinyrouter.protocol import CurveProtocol, ProtocolError +from tinyrouter.runs import run_one + +pytestmark = pytest.mark.filterwarnings("ignore::tinyrouter.calibrate.TemperatureBoundWarning") + +ROOT = Path(__file__).parent.parent +BERT_NAME = "google-bert/bert-base-uncased" +MODERN_NAME = "answerdotai/ModernBERT-base" + + +def make_protocol(tmp_path, s_min=None, bert_lr=None, modern_lr=None, **bert_changes): + def here(config): + return replace( + config, checkpoint_root=str(tmp_path / "ckpt"), results_root=str(tmp_path / "res") + ) + + bert = here(replace(load_config(ROOT / "configs" / "bert-base.yaml"), **bert_changes)) + modern = here(load_config(ROOT / "configs" / "modernbert-base.yaml")) + return CurveProtocol( + s_min, + {"bert": bert_lr, "modernbert": modern_lr}, + {"bert": bert, "modernbert": modern}, + {"bert": "b", "modernbert": "m"}, + ) + + +def quiet(_: str) -> None: + pass + + +def keys_anywhere(value: object) -> set[str]: + if isinstance(value, dict): + return set(value) | set().union(*(keys_anywhere(v) for v in value.values())) + if isinstance(value, list): + return set().union(*(keys_anywhere(v) for v in value)) if value else set() + return set() + + +def test_scoring_refuses_test_logits(): + with pytest.raises(LeakageError, match="only see validation"): + pilots.validation_scores(fake_splits()["test"]) + + +def test_prediction_refuses_the_test_split_before_touching_a_model(tmp_path): + split = Split("test", ("a",), np.array([0])) + with pytest.raises(LeakageError): + pilots.predict_validation( + tmp_path / "no-such-model", split, load_config(ROOT / "configs/bert-base.yaml") + ) + + +def test_validation_scores_count_in_scope_and_oos_hits(): + scores = pilots.validation_scores(fake_validation(2_901, oos_correct=68)) + assert scores["in_scope_correct"] == 2_901 and scores["in_scope_n"] == 3_000 + assert scores["oos_correct"] == 68 and scores["oos_n"] == 100 + assert scores["in_scope_accuracy_150"] == 2_901 / 3_000 + + +def entry(model, value, in_scope, oos=0): + return { + "model": model, + "value": value, + "validation": pilots.validation_scores(fake_validation(in_scope, oos)), + } + + +def test_lr_selection_breaks_ties_on_oos_recall_then_on_the_smaller_rate(): + entries = [ + entry("bert", 1e-5, 2_900, 60), + entry("bert", 2e-5, 2_900, 70), + entry("bert", 5e-5, 2_899, 99), + entry("modernbert", 1e-5, 2_950, 50), + entry("modernbert", 2e-5, 2_950, 50), + entry("modernbert", 5e-5, 2_949, 90), + ] + assert pilots.select_lr(entries) == {"bert": 2e-5, "modernbert": 1e-5} + + +def test_steps_selection_uses_both_encoders_and_prefers_the_smaller_s_min_on_a_tie(): + entries = [ + entry("bert", 100, 1_000), + entry("modernbert", 100, 1_300), + entry("bert", 200, 1_200), + entry("modernbert", 200, 1_100), + entry("bert", 400, 1_500), + entry("modernbert", 400, 700), + ] + assert pilots.select_steps(entries) == 100 # 2,300 each for 100 and 200 + entries[2]["validation"]["in_scope_correct"] += 1 + assert pilots.select_steps(entries) == 200 + + +def ac2_donor_run(protocol: CurveProtocol, in_scope_correct: int = 0) -> None: + fake = FakeRuns() + fake.correct = {(BERT_NAME, 5e-5): in_scope_correct} + run_one( + protocol.base("bert").with_seed(42), fake.train, fake.evaluate, quiet, keep_weights=True + ) + + +def test_pilot_lr_reuses_ac2_for_bert_5e5_and_trains_the_other_five(tmp_path): + protocol = make_protocol(tmp_path) + ac2_donor_run(protocol, in_scope_correct=2_902) + fake = FakeRuns() + fake.correct = {(MODERN_NAME, 2e-5): 2_950} + out = tmp_path / "res" / "pilots" / "lr.json" + body = pilots.run_pilot("lr", protocol, out, fake.train, fake.validation, quiet) + assert len(fake.trained) == 5 + assert all("lr5e-05" not in t or "ModernBERT" in t for t in fake.trained) + reused = [p for p in body["points"] if p["source"].startswith("reused")] + assert [(p["model"], p["value"]) for p in reused] == [("bert", 5e-5)] + assert body["selected"] == {"bert": 5e-5, "modernbert": 2e-5} + assert json.loads(out.read_text()) == json.loads(json.dumps(body)) + # The AC2 point's numbers come from its archive's validation logits. + ac2 = next(p for p in body["points"] if p["source"].startswith("reused")) + assert ac2["validation"]["in_scope_correct"] == 2_902 + assert ac2["training"]["global_step"] == 2385 + + +def test_pilot_output_holds_no_test_numbers_at_all(tmp_path): + protocol = make_protocol(tmp_path) + ac2_donor_run(protocol) + fake = FakeRuns() + body = pilots.run_pilot( + "lr", protocol, tmp_path / "lr.json", fake.train, fake.validation, quiet + ) + assert not {k for k in keys_anywhere(body) if "test" in k} + assert body["split"] == "validation" + + +def test_pilot_never_calls_the_full_evaluation(tmp_path, monkeypatch): + import tinyrouter.evaluate as evaluate_module + + def forbidden(*_): + raise AssertionError("pilots must not run evaluate(), which scores test") + + monkeypatch.setattr(evaluate_module, "evaluate", forbidden) + monkeypatch.setattr(evaluate_module, "score", forbidden) + protocol = make_protocol(tmp_path) + fake = FakeRuns() + pilots.run_pilot("lr", protocol, tmp_path / "lr.json", fake.train, fake.validation, quiet) + + +def test_an_ac2_run_made_with_other_settings_is_not_reused(tmp_path): + ac2_donor_run(make_protocol(tmp_path, warmup_ratio=0.0)) + protocol = make_protocol(tmp_path) # the config file now says warmup 0.1 + fake = FakeRuns() + body = pilots.run_pilot( + "lr", protocol, tmp_path / "lr.json", fake.train, fake.validation, quiet + ) + assert len(fake.trained) == 6 + assert {p["source"] for p in body["points"]} == {"trained"} + + +def test_pilot_weights_are_deleted_and_ac2_weights_are_left_alone(tmp_path): + protocol = make_protocol(tmp_path) + ac2_donor_run(protocol) + fake = FakeRuns() + pilots.run_pilot("lr", protocol, tmp_path / "lr.json", fake.train, fake.validation, quiet) + ckpt = tmp_path / "ckpt" + assert not list((ckpt / "pilots").rglob("model.safetensors")) + assert (ckpt / "bert-base-uncased-full-seed42" / "final" / "model.safetensors").exists() + + +def test_rerunning_a_finished_pilot_trains_nothing(tmp_path): + protocol = make_protocol(tmp_path) + out = tmp_path / "lr.json" + pilots.run_pilot("lr", protocol, out, FakeRuns().train, FakeRuns().validation, quiet) + again = FakeRuns() + pilots.run_pilot("lr", protocol, out, again.train, again.validation, quiet) + assert again.trained == [] + + +def test_pilot_steps_needs_the_learning_rates_first(tmp_path): + with pytest.raises(ProtocolError, match="pilot-lr"): + pilots.run_pilot("steps", make_protocol(tmp_path), tmp_path / "s.json") + + +def test_pilot_steps_trains_k5_for_each_s_min_and_encoder(tmp_path): + protocol = make_protocol(tmp_path, bert_lr=5e-5, modern_lr=2e-5) + fake = FakeRuns() + fake.correct = {(BERT_NAME, 200): 700, (MODERN_NAME, 200): 800, (MODERN_NAME, 400): 1_600} + body = pilots.run_pilot( + "steps", protocol, tmp_path / "s.json", fake.train, fake.validation, quiet + ) + assert len(fake.trained) == 6 + assert body["selected"] == 400 + steps = {(p["model"], p["value"]): p["training"]["step_plan"] for p in body["points"]} + assert steps[("bert", 100)]["decided_by"] == "epochs" # 120 epoch steps at k=5 + assert steps[("bert", 400)]["decided_by"] == "min_train_steps" + assert {p["config"]["k_shot"] for p in body["points"]} == {5} + + +def test_a_crash_after_training_keeps_the_weights_for_the_rerun(tmp_path): + protocol = make_protocol(tmp_path) + fake = FakeRuns() + + def broken(config, model_dir) -> SplitLogits: + raise RuntimeError("prediction crashed") + + with pytest.raises(RuntimeError, match="crashed"): + pilots.run_pilot("lr", protocol, tmp_path / "lr.json", fake.train, broken, quiet) + again = FakeRuns() + pilots.run_pilot("lr", protocol, tmp_path / "lr.json", again.train, again.validation, quiet) + assert len(again.trained) == len(pilots.lr_points(protocol)) - 1 diff --git a/tests/test_protocol.py b/tests/test_protocol.py new file mode 100644 index 0000000..c89f1f0 --- /dev/null +++ b/tests/test_protocol.py @@ -0,0 +1,113 @@ +from dataclasses import asdict, replace +from pathlib import Path + +import pytest + +from tinyrouter.config import load_config +from tinyrouter.protocol import ( + LR_GRID, + S_MIN_GRID, + SEEDS, + CurveProtocol, + ProtocolError, + ac2_donor, + load_protocol, +) +from tinyrouter.sampling import CURVE_KS + +ROOT = Path(__file__).parent.parent + + +def protocol_file(tmp_path, s_min="null", bert="null", modern="null", extra="") -> Path: + path = tmp_path / "curve.yaml" + path.write_text( + f"min_train_steps: {s_min}\nmodels:\n" + f" bert:\n config: {ROOT}/configs/bert-base.yaml\n learning_rate: {bert}\n" + f" modernbert:\n config: {ROOT}/configs/modernbert-base.yaml\n" + f" learning_rate: {modern}\n{extra}" + ) + return path + + +def decided(tmp_path) -> CurveProtocol: + return load_protocol(protocol_file(tmp_path, s_min=400, bert="5.0e-5", modern="2.0e-5")) + + +def test_the_committed_protocol_loads_with_both_pilots_still_open(): + protocol = load_protocol(ROOT / "configs" / "curve.yaml") + assert protocol.min_train_steps is None + assert protocol.learning_rates == {"bert": None, "modernbert": None} + with pytest.raises(ProtocolError, match="make pilot-lr"): + protocol.curve_configs("bert") + + +def test_curve_refuses_to_start_until_s_min_is_chosen(tmp_path): + protocol = load_protocol(protocol_file(tmp_path, bert="5.0e-5", modern="2.0e-5")) + with pytest.raises(ProtocolError, match="make pilot-steps"): + protocol.curve_configs("modernbert") + + +@pytest.mark.parametrize( + "fields", + [ + {"bert": "3.0e-5"}, + {"bert": "5e-5"}, # YAML reads this as a string + {"modern": "true"}, + {"s_min": 150}, + {"s_min": "100.0"}, + {"extra": " extra_model:\n config: x\n learning_rate: null\n"}, + ], +) +def test_values_off_the_protocol_grid_are_refused(tmp_path, fields): + with pytest.raises(ProtocolError): + load_protocol(protocol_file(tmp_path, **fields)) + + +def test_curve_configs_cover_every_k_and_seed_with_one_lr_and_one_s_min(tmp_path): + protocol = decided(tmp_path) + configs = protocol.curve_configs("modernbert") + assert [(c.k_shot, c.seed) for c in configs] == [(k, s) for k in CURVE_KS for s in SEEDS] + assert {c.learning_rate for c in configs} == {2e-5} + assert {c.min_train_steps for c in configs} == {400} + assert {c.oos_train for c in configs} == {None} + + +def test_a_curve_run_differs_from_its_base_config_only_in_lr_s_min_k_and_seed(tmp_path): + base = load_config(ROOT / "configs" / "bert-base.yaml") + point = decided(tmp_path).curve_config("bert", 25, 43) + changed = {k for k, v in asdict(point).items() if asdict(base)[k] != v} + assert changed == {"min_train_steps", "k_shot", "seed"} # bert's lr 5e-5 equals the file's + + +def test_the_two_encoders_differ_only_in_model_and_learning_rate(): + bert = asdict(load_config(ROOT / "configs" / "bert-base.yaml")) + modern = asdict(load_config(ROOT / "configs" / "modernbert-base.yaml")) + assert {k for k in bert if bert[k] != modern[k]} == { + "model_name", + "model_revision", + "learning_rate", + } + + +def test_the_ablation_is_modernbert_at_k100_with_no_oos_rows(tmp_path): + configs = decided(tmp_path).ablation_configs() + assert [c.seed for c in configs] == list(SEEDS) + assert all(c.k_shot == 100 and c.oos_train == 0 for c in configs) + assert {c.model_name for c in configs} == {"answerdotai/ModernBERT-base"} + assert configs[0].run_name == "ModernBERT-base-k100-oos0-seed42" + + +def test_only_bert_at_k100_with_the_table_oos_count_has_an_ac2_donor(tmp_path): + protocol = decided(tmp_path) + donor = ac2_donor(protocol, protocol.curve_config("bert", 100, 44)) + assert donor == protocol.base("bert").with_seed(44) + assert donor.run_name == "bert-base-uncased-full-seed44" + assert ac2_donor(protocol, protocol.curve_config("bert", 50, 44)) is None + assert ac2_donor(protocol, protocol.curve_config("modernbert", 100, 44)) is None + ablation_like = replace(protocol.curve_config("bert", 100, 44), oos_train=0) + assert ac2_donor(protocol, ablation_like) is None + + +def test_grids_are_the_ones_the_plan_fixes(): + assert LR_GRID == (1e-5, 2e-5, 5e-5) + assert S_MIN_GRID == (100, 200, 400) diff --git a/tests/test_runs.py b/tests/test_runs.py new file mode 100644 index 0000000..a9871c1 --- /dev/null +++ b/tests/test_runs.py @@ -0,0 +1,142 @@ +"""runs.equivalent_run: when an existing run may stand in for a curve point, and when not.""" + +import json +from dataclasses import asdict, replace +from pathlib import Path + +import pytest + +from run_fakes import FakeRuns +from tinyrouter.config import ADDED_FIELD_DEFAULTS, load_config +from tinyrouter.evaluate import RunPaths +from tinyrouter.runs import equivalent_run, reusable_equivalent, run_one + +pytestmark = pytest.mark.filterwarnings("ignore::tinyrouter.calibrate.TemperatureBoundWarning") + +ROOT = Path(__file__).parent.parent +BERT = load_config(ROOT / "configs" / "bert-base.yaml") +AC2_RECORDS = { + s: ROOT / "results" / "runs" / f"bert-base-uncased-full-seed{s}.json" for s in (42, 43, 44) +} + + +def curve_point(seed: int = 42, **fields: object): + return replace(BERT, **{"k_shot": 100, "min_train_steps": 400, "seed": seed, **fields}) + + +def ac2_like_record(seed: int = 42, training: dict | None = None) -> dict: + """The shape of an AC2 results JSON: a config written before k_shot and S_min existed.""" + config = { + k: v for k, v in asdict(BERT.with_seed(seed)).items() if k not in ADDED_FIELD_DEFAULTS + } + return { + "run_name": BERT.with_seed(seed).run_name, + "config": config, + "training": { + "seed": seed, + "train_rows": 15_250, + "oos_train_rows": 250, + "global_step": 2385, + **(training or {}), + }, + } + + +@pytest.mark.parametrize("seed", [42, 43, 44]) +def test_the_committed_ac2_runs_are_equivalent_to_bert_k100_at_lr_5e5_for_every_s_min(seed): + record = json.loads(AC2_RECORDS[seed].read_text()) + assert "k_shot" not in record["config"] # written before the field existed + for s_min in (None, 100, 200, 400): + assert equivalent_run(curve_point(seed, min_train_steps=s_min), record) + + +def test_an_ac2_run_is_not_equivalent_to_another_seed(): + record = json.loads(AC2_RECORDS[42].read_text()) + assert not equivalent_run(curve_point(43), record) + + +@pytest.mark.parametrize( + "change", + [ + {"learning_rate": 2e-5}, + {"learning_rate": 1e-5}, + {"model_revision": "0" * 40}, + {"max_length": 128}, + {"warmup_ratio": 0.0}, + {"weight_decay": 0.0}, + {"train_batch_size": 16}, + {"num_train_epochs": 3.0}, + {"eval_batch_size": 64}, + {"device": "cpu"}, + {"k_shot": 50}, + {"oos_train": 0}, + ], + ids=lambda c: next(iter(c)), +) +def test_any_other_difference_means_not_the_same_run(change): + assert equivalent_run(curve_point(), ac2_like_record()) + assert not equivalent_run(curve_point(**change), ac2_like_record()) + + +@pytest.mark.parametrize( + "training", + [{"global_step": 2384}, {"train_rows": 15_249}, {"oos_train_rows": 100}, {"seed": 43}], + ids=lambda t: next(iter(t)), +) +def test_a_record_whose_own_training_summary_disagrees_is_refused(training): + record = ac2_like_record(training=training) + assert record["config"]["seed"] == 42 + assert not equivalent_run(curve_point(), record) + + +def test_a_record_missing_a_field_that_is_not_a_known_later_addition_is_refused(): + record = ac2_like_record() + del record["config"]["warmup_ratio"] + assert not equivalent_run(curve_point(), record) + + +def test_a_record_carrying_an_unknown_field_is_refused(): + record = ac2_like_record() + record["config"]["label_smoothing"] = 0.1 + assert not equivalent_run(curve_point(), record) + + +def test_min_train_steps_that_decides_equals_the_same_fixed_max_steps(): + """Rewrite 2: 200 steps via S_min at k=5 is the same TrainingArguments as max_steps=200.""" + fixed = replace(BERT, k_shot=5, max_steps=200) + by_floor = replace(BERT, k_shot=5, min_train_steps=200) + record = { + "run_name": fixed.run_name, + "config": asdict(fixed), + "training": {"seed": 42, "train_rows": 763, "oos_train_rows": 13, "global_step": 200}, + } + assert equivalent_run(by_floor, record) + assert not equivalent_run(replace(by_floor, min_train_steps=400), record) + assert not equivalent_run(replace(BERT, k_shot=5), record) # 120 epoch steps, not 200 + + +def test_per_intent_cap_is_never_treated_as_a_k_shot_sample(): + record = ac2_like_record() + record["config"]["per_intent"] = 100 + assert not equivalent_run(curve_point(), record) + + +@pytest.fixture +def tmp_bert(tmp_path): + return replace(BERT, checkpoint_root=str(tmp_path / "ckpt"), results_root=str(tmp_path / "res")) + + +def test_reusable_equivalent_needs_an_intact_archive(tmp_bert): + donor = tmp_bert.with_seed(42) + fake = FakeRuns() + run_one(donor, fake.train, fake.evaluate, lambda _: None, keep_weights=False) + point = replace(tmp_bert, k_shot=100, min_train_steps=400) + assert reusable_equivalent(point, donor) is not None + archive = RunPaths.of(donor).logits + archive.write_bytes(archive.read_bytes() + b"x") + assert reusable_equivalent(point, donor) is None + + +def test_reusable_equivalent_is_none_when_the_donor_never_ran(tmp_bert): + point = replace(tmp_bert, k_shot=100, min_train_steps=400) + assert reusable_equivalent(point, tmp_bert) is None diff --git a/tests/test_sampling.py b/tests/test_sampling.py new file mode 100644 index 0000000..188e0e1 --- /dev/null +++ b/tests/test_sampling.py @@ -0,0 +1,126 @@ +import math +from collections import Counter + +import numpy as np +import pytest + +from tinyrouter.data import Split +from tinyrouter.labels import load_label_space +from tinyrouter.sampling import ( + CURVE_KS, + OOS_TRAIN_ROWS, + SamplingError, + curve_sample, + planned_rows, + sample_fingerprint, + sample_k_shot, +) + +OOS = load_label_space().oos_intent_id + + +def full_train() -> Split: + """Shaped like the OOS+ train split: 100 rows per in-scope intent, 250 oos, interleaved.""" + rng = np.random.default_rng(0) + intents = np.array([i for i in range(151) if i != OOS for _ in range(100)] + [OOS] * 250) + intents = intents[rng.permutation(len(intents))] + texts = tuple(f"row{n} intent{i}" for n, i in enumerate(intents)) + return Split("train", texts, intents.astype(np.int64)) + + +FULL = full_train() + + +def test_table_is_ceil_of_two_and_a_half_k(): + assert {k: math.ceil(2.5 * k) for k in CURVE_KS} == OOS_TRAIN_ROWS + assert CURVE_KS == (1, 5, 10, 25, 50, 100) + + +@pytest.mark.parametrize("k", CURVE_KS) +def test_every_in_scope_intent_gets_exactly_k_rows_and_oos_the_table_count(k): + counts = Counter(sample_k_shot(FULL, k, seed=42).intents.tolist()) + assert counts.pop(OOS) == OOS_TRAIN_ROWS[k] + assert len(counts) == 150 and set(counts.values()) == {k} + + +@pytest.mark.parametrize("k", CURVE_KS) +def test_planned_rows_match_the_drawn_sample(k): + sample = sample_k_shot(FULL, k, seed=7) + assert planned_rows(k) == (len(sample), int((sample.intents == OOS).sum())) + + +def test_k_100_is_the_whole_train_split_in_its_original_order(): + sample = sample_k_shot(FULL, 100, seed=43) + assert Counter(zip(sample.texts, sample.intents.tolist(), strict=True)) == Counter( + zip(FULL.texts, FULL.intents.tolist(), strict=True) + ) + assert sample.texts == FULL.texts + assert sample_fingerprint(sample) == sample_fingerprint(FULL) + + +@pytest.mark.parametrize("k", [1, 5, 10, 25, 50]) +def test_same_seed_draws_the_same_rows_and_another_seed_different_ones(k): + a, b = sample_k_shot(FULL, k, seed=42), sample_k_shot(FULL, k, seed=42) + assert a.texts == b.texts + assert sample_k_shot(FULL, k, seed=43).texts != a.texts + + +def test_selected_rows_keep_their_original_order(): + sample = sample_k_shot(FULL, 5, seed=42) + positions = [FULL.texts.index(t) for t in sample.texts] + assert positions == sorted(positions) + + +@pytest.mark.parametrize("k", [0, 2, 20, 99, 101, -1]) +def test_k_not_in_the_table_is_refused(k): + with pytest.raises(SamplingError, match="not a curve point"): + sample_k_shot(FULL, k, seed=42) + + +def test_oos_override_zero_drops_every_oos_row_for_the_ablation(): + sample = sample_k_shot(FULL, 100, seed=42, oos_override=0) + assert OOS not in sample.intents.tolist() + assert len(sample) == 15_000 + assert planned_rows(100, 0) == (15_000, 0) + + +def test_negative_oos_override_is_refused(): + with pytest.raises(SamplingError, match=">= 0"): + sample_k_shot(FULL, 5, seed=42, oos_override=-1) + + +def test_only_the_train_split_is_sampled(): + with pytest.raises(SamplingError, match="only the train split"): + sample_k_shot(Split("test", FULL.texts, FULL.intents), 5, seed=42) + + +def test_a_split_missing_an_intent_is_refused(): + keep = np.flatnonzero(FULL.intents != 3) + with pytest.raises(SamplingError, match="distinct intents"): + sample_k_shot(FULL.take(keep), 5, seed=42) + + +def test_an_intent_with_too_few_rows_is_refused(): + rows_of_7 = np.flatnonzero(FULL.intents == 7) + keep = np.setdiff1d(np.arange(len(FULL)), rows_of_7[:95]) + with pytest.raises(SamplingError, match="intent 7 has 5 rows, need 10"): + sample_k_shot(FULL.take(keep), 10, seed=42) + + +def test_curve_sample_reads_the_train_split(monkeypatch): + import tinyrouter.sampling as sampling + + asked: list[str] = [] + monkeypatch.setattr(sampling, "load_split", lambda name: asked.append(name) or FULL) + assert curve_sample(5, 42).texts == sample_k_shot(FULL, 5, 42).texts + assert asked == ["train"] + + +@pytest.mark.network +def test_on_the_real_train_split_k_100_is_everything_and_k_5_is_balanced(): + from tinyrouter.data import load_split + + real = load_split("train") + assert sample_k_shot(real, 100, seed=42).texts == real.texts + counts = Counter(sample_k_shot(real, 5, seed=42).intents.tolist()) + assert counts.pop(OOS) == 13 and set(counts.values()) == {5} diff --git a/tests/test_steps.py b/tests/test_steps.py new file mode 100644 index 0000000..1139873 --- /dev/null +++ b/tests/test_steps.py @@ -0,0 +1,64 @@ +from dataclasses import replace + +import pytest + +from tinyrouter.config import RunConfig +from tinyrouter.steps import epoch_steps, plan_steps + +BASE = RunConfig(model_name="m", model_revision="r") + + +def test_epoch_steps_follow_the_trainer_arithmetic(): + assert epoch_steps(15_250, 32, 5) == 2385 # the AC2 run's recorded global_step + assert epoch_steps(763, 32, 5) == 120 # k=5: 24 batches per epoch + assert epoch_steps(153, 32, 5) == 25 # k=1 + assert epoch_steps(33, 32, 1.5) == 3 # ceil(1.5 * 2) + + +def test_empty_training_set_cannot_be_planned(): + with pytest.raises(ValueError, match="empty"): + epoch_steps(0, 32, 5) + + +def test_without_min_train_steps_the_epochs_decide(): + plan = plan_steps(BASE, 763) + assert (plan.planned_steps, plan.decided_by, plan.max_steps_arg) == (120, "epochs", -1) + + +@pytest.mark.parametrize(("s_min", "rows", "steps"), [(400, 153, 400), (200, 763, 200)]) +def test_min_train_steps_above_the_epoch_count_sets_the_steps(s_min, rows, steps): + plan = plan_steps(replace(BASE, min_train_steps=s_min), rows) + assert plan.planned_steps == plan.max_steps_arg == steps + assert plan.decided_by == "min_train_steps" + assert plan.epoch_steps == epoch_steps(rows, 32, 5) + + +@pytest.mark.parametrize("s_min", [100, 120]) +def test_min_train_steps_at_or_below_the_epoch_count_changes_nothing(s_min): + plan = plan_steps(replace(BASE, min_train_steps=s_min), 763) + assert (plan.planned_steps, plan.decided_by, plan.max_steps_arg) == (120, "epochs", -1) + assert plan.min_train_steps == s_min + + +def test_full_data_is_never_extended_by_any_s_min_on_the_grid(): + for s_min in (100, 200, 400): + assert plan_steps(replace(BASE, min_train_steps=s_min), 15_250).decided_by == "epochs" + + +def test_an_explicit_max_steps_wins(): + plan = plan_steps(replace(BASE, max_steps=7), 15_250) + assert (plan.planned_steps, plan.decided_by, plan.max_steps_arg) == (7, "max_steps", 7) + + +@pytest.mark.parametrize( + ("fields", "message"), + [ + ({"min_train_steps": 0}, ">= 1"), + ({"min_train_steps": 10, "max_steps": 5}, "contradict"), + ({"k_shot": 5, "per_intent": 5}, "not both"), + ({"oos_train": 0}, "needs k_shot"), + ], +) +def test_contradictory_configs_are_refused(fields, message): + with pytest.raises(ValueError, match=message): + replace(BASE, **fields) diff --git a/tests/test_train_smoke.py b/tests/test_train_smoke.py index b32a823..de94d02 100644 --- a/tests/test_train_smoke.py +++ b/tests/test_train_smoke.py @@ -276,3 +276,40 @@ def boom(*_): with pytest.raises(RuntimeError, match="prediction failed"): evaluate(config, model_dir) assert not old.exists() + + +@pytest.mark.parametrize( + ("fields", "steps", "decided_by"), + [ + ({"num_train_epochs": 1.0}, 4, "epochs"), # 32 rows / batch 8 + ({"num_train_epochs": 1.0, "min_train_steps": 6}, 6, "min_train_steps"), + ({"num_train_epochs": 2.0, "min_train_steps": 3}, 8, "epochs"), + ], +) +def test_trainer_runs_exactly_the_planned_steps_and_records_the_plan( + tiny_model_dir, tmp_path, fields, steps, decided_by +): + import json + + config = replace(smoke_config(tiny_model_dir, tmp_path), max_steps=-1, **fields) + final_dir = train(config, synthetic_split("train"), tmp_path / "ckpt" / config.run_name) + summary = json.loads((final_dir / "train_summary.json").read_text()) + assert summary["global_step"] == summary["step_plan"]["planned_steps"] == steps + assert summary["step_plan"]["decided_by"] == decided_by + assert list((tmp_path / "ckpt" / config.run_name).glob("checkpoint-*")) == [] + + +def test_step_count_that_differs_from_the_plan_raises(tiny_model_dir, tmp_path, monkeypatch): + import tinyrouter.train as train_module + from tinyrouter.steps import StepPlan + + config = replace(smoke_config(tiny_model_dir, tmp_path), max_steps=-1, num_train_epochs=1.0) + monkeypatch.setattr(train_module, "plan_steps", lambda *_: StepPlan(4, None, 5, "epochs", -1)) + with pytest.raises(train_module.StepCountError, match="trained 4 steps, planned 5"): + train(config, synthetic_split("train"), tmp_path / "x") + + +def test_training_refuses_a_learning_rate_still_waiting_for_its_pilot(tmp_path): + config = replace(smoke_config(Path("unused"), tmp_path), learning_rate=None) + with pytest.raises(ValueError, match="make pilot-lr"): + train(config, synthetic_split("train"), tmp_path / "x") From 298c87972c1653a3b29a7a6f325038e615947a01 Mon Sep 17 00:00:00 2001 From: drewOrc <36374426+drewOrc@users.noreply.github.com> Date: Wed, 23 Sep 2026 14:44:16 +0800 Subject: [PATCH 2/4] Add ModernBERT compatibility trial and ONNX export check 50 steps on MPS for ModernBERT-base and bert-base-uncased with the same batch, max_length and step count, then one validation pass; and an ONNX export of the ModernBERT trial model compared with PyTorch on CPU. Results in results/compat/. ONNX tooling is an optional dependency group (uv sync --group onnx), not installed in CI. --- pyproject.toml | 7 + results/compat/bert-base-trial.json | 56 ++++++++ results/compat/modernbert-base-onnx.json | 42 ++++++ results/compat/modernbert-base-trial.json | 56 ++++++++ scripts/compat_trial.py | 147 ++++++++++++++++++++ scripts/export_onnx.py | 155 ++++++++++++++++++++++ uv.lock | 121 +++++++++++++++++ 7 files changed, 584 insertions(+) create mode 100644 results/compat/bert-base-trial.json create mode 100644 results/compat/modernbert-base-onnx.json create mode 100644 results/compat/modernbert-base-trial.json create mode 100644 scripts/compat_trial.py create mode 100644 scripts/export_onnx.py diff --git a/pyproject.toml b/pyproject.toml index 3ade08b..4c2cb96 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -37,6 +37,13 @@ dev = [ "pytest==9.1.1", "ruff==0.16.8", ] +# ONNX export and ONNX Runtime (RQ7; scripts/export_onnx.py). Not installed +# by `uv sync` or in CI; `uv sync --group onnx` installs it. +onnx = [ + "onnx==1.23.0", + "onnxruntime==1.30.0", + "onnxscript==0.7.2", +] [tool.uv] default-groups = ["dev"] diff --git a/results/compat/bert-base-trial.json b/results/compat/bert-base-trial.json new file mode 100644 index 0000000..d4afe50 --- /dev/null +++ b/results/compat/bert-base-trial.json @@ -0,0 +1,56 @@ +{ + "config": "configs/bert-base.yaml", + "model_name": "google-bert/bert-base-uncased", + "model_revision": "86b5e0934494bd15c9632b12f734a8a67f723594", + "learning_rate": 5e-05, + "train_batch_size": 32, + "max_length": 64, + "steps": 50, + "wall_seconds": 18.426, + "step_seconds_median": 0.139, + "step_seconds_mean_after_5": 0.1445, + "step_seconds_first": 0.5056, + "estimated_full_5_epoch_seconds": 331.6, + "loss_first_5": [ + 5.080526828765869, + 5.005545616149902, + 5.043530464172363, + 5.185548305511475, + 5.147557735443115 + ], + "loss_last_5": [ + 4.936195373535156, + 4.977850914001465, + 4.936770439147949, + 4.952746391296387, + 5.00227165222168 + ], + "loss_decreased": true, + "any_nan_loss": false, + "validation_logits_finite": true, + "validation_in_scope_accuracy_150": 0.021666666666666667, + "validation_oos_recall_151": 0.0, + "parameters": { + "total": 109598359, + "trainable": 109598359 + }, + "attn_implementation": "sdpa", + "peak_memory": { + "device": "mps", + "measures": "sampled after each backward and optimizer step: torch.mps.driver_allocated_memory (Metal driver, includes allocator cache; headline, usually at or above the tensor peak but not guaranteed) and current_allocated_memory (live tensors; lower bound)", + "samples": 101, + "driver_allocated_bytes": 3414933504, + "tensor_allocated_bytes": 1780597248 + }, + "environment": { + "python": "3.12.12", + "torch": "2.14.0", + "transformers": "5.17.0", + "platform": "macOS-26.6.2-arm64-arm-64bit", + "device": "mps", + "device_name": "Apple M4", + "deterministic_algorithms": false, + "dataset_revision": "155b9c710419136e17307b80d0a13e68cd46b4ec" + }, + "plan_note": "full 5 epochs = StepPlan(epoch_steps=2385, min_train_steps=None, planned_steps=2385, decided_by='epochs', max_steps_arg=-1)" +} diff --git a/results/compat/modernbert-base-onnx.json b/results/compat/modernbert-base-onnx.json new file mode 100644 index 0000000..da15df8 --- /dev/null +++ b/results/compat/modernbert-base-onnx.json @@ -0,0 +1,42 @@ +{ + "model_type": "modernbert", + "rows": 10, + "padded_sequence_length": 14, + "attempts": [ + { + "exporter": "dynamo", + "ok": true, + "seconds": 14.6, + "onnx_bytes": 601513367, + "max_abs_logit_diff": 2.1219253540039062e-05, + "other_shape": [ + 3, + 25 + ], + "other_shape_max_abs_logit_diff": 1.1337921023368835e-05, + "argmax_agree": 10, + "rows": 10 + }, + { + "exporter": "torchscript", + "ok": true, + "seconds": 3.6, + "onnx_bytes": 599490381, + "max_abs_logit_diff": 2.1219253540039062e-05, + "other_shape": [ + 3, + 25 + ], + "other_shape_max_abs_logit_diff": 1.1337921023368835e-05, + "argmax_agree": 10, + "rows": 10 + } + ], + "versions": { + "torch": "2.14.0", + "transformers": "5.17.0", + "onnx": "1.23.0", + "onnxruntime": "1.30.0" + }, + "machine": "macOS-26.6.2-arm64-arm-64bit" +} diff --git a/results/compat/modernbert-base-trial.json b/results/compat/modernbert-base-trial.json new file mode 100644 index 0000000..b73733f --- /dev/null +++ b/results/compat/modernbert-base-trial.json @@ -0,0 +1,56 @@ +{ + "config": "configs/modernbert-base.yaml", + "model_name": "answerdotai/ModernBERT-base", + "model_revision": "8949b909ec900327062f0ebf497f51aef5e6f0c8", + "learning_rate": 5e-05, + "train_batch_size": 32, + "max_length": 64, + "steps": 50, + "wall_seconds": 30.622, + "step_seconds_median": 0.2344, + "step_seconds_mean_after_5": 0.2778, + "step_seconds_first": 1.1259, + "estimated_full_5_epoch_seconds": 559.1, + "loss_first_5": [ + 5.222139358520508, + 5.159377098083496, + 5.297676086425781, + 5.101874351501465, + 5.049973964691162 + ], + "loss_last_5": [ + 4.463723182678223, + 4.428763389587402, + 4.537565231323242, + 4.791100025177002, + 4.218628883361816 + ], + "loss_decreased": true, + "any_nan_loss": false, + "validation_logits_finite": true, + "validation_in_scope_accuracy_150": 0.11266666666666666, + "validation_oos_recall_151": 0.1, + "parameters": { + "total": 149720983, + "trainable": 149720983 + }, + "attn_implementation": "sdpa", + "peak_memory": { + "device": "mps", + "measures": "sampled after each backward and optimizer step: torch.mps.driver_allocated_memory (Metal driver, includes allocator cache; headline, usually at or above the tensor peak but not guaranteed) and current_allocated_memory (live tensors; lower bound)", + "samples": 101, + "driver_allocated_bytes": 5151735808, + "tensor_allocated_bytes": 2464530176 + }, + "environment": { + "python": "3.12.12", + "torch": "2.14.0", + "transformers": "5.17.0", + "platform": "macOS-26.6.2-arm64-arm-64bit", + "device": "mps", + "device_name": "Apple M4", + "deterministic_algorithms": false, + "dataset_revision": "155b9c710419136e17307b80d0a13e68cd46b4ec" + }, + "plan_note": "full 5 epochs = StepPlan(epoch_steps=2385, min_train_steps=None, planned_steps=2385, decided_by='epochs', max_steps_arg=-1)" +} diff --git a/scripts/compat_trial.py b/scripts/compat_trial.py new file mode 100644 index 0000000..4e2ed90 --- /dev/null +++ b/scripts/compat_trial.py @@ -0,0 +1,147 @@ +"""Short compatibility trial of an encoder on this machine (docs/PLAN.md section 4). + +Trains ``--steps`` optimizer steps on the full train split with the +config's batch size and max_length, then scores validation once. Records +the wall time of every step, the sampled peak memory, the loss of every +step and whether any loss or logit was NaN. Used to check that +ModernBERT trains on MPS and to compare its per-step time with +bert-base-uncased under identical settings. + +The numbers after 50 steps are not a result; the model is barely trained. +Validation only: this script never loads the test split. + + uv run python scripts/compat_trial.py --config configs/modernbert-base.yaml \\ + --learning-rate 5e-5 --out results/compat/modernbert-base.json --keep-weights DIR +""" + +from __future__ import annotations + +import argparse +import json +import math +import shutil +import time +from dataclasses import replace +from pathlib import Path + +import numpy as np + +from tinyrouter.config import load_config +from tinyrouter.data import load_split +from tinyrouter.efficiency import ( + count_parameters, + memory_report, + sampler_callback, + start_measurement, +) +from tinyrouter.evaluate import environment, predict_logits +from tinyrouter.labels import load_label_space +from tinyrouter.metrics import in_scope_accuracy, oos_recall +from tinyrouter.steps import plan_steps +from tinyrouter.train import ( + load_model_and_tokenizer, + pick_device, + to_hf_dataset, + training_arguments, +) + +AC2_EPOCH_STEPS = 2385 + + +def step_timer(times: list[float]) -> object: + from transformers import TrainerCallback + + class StepTimer(TrainerCallback): + def on_step_begin(self, args, state, control, **kwargs): # noqa: ANN001 + self.started = time.perf_counter() + + def on_step_end(self, args, state, control, **kwargs): # noqa: ANN001 + times.append(time.perf_counter() - self.started) + + return StepTimer() + + +def run_trial(config_path: str, steps: int, lr: float, work: Path) -> tuple[dict, Path]: + from transformers import DataCollatorWithPadding, Trainer, set_seed + + config = replace( + load_config(config_path), learning_rate=lr, max_steps=steps, checkpoint_root=str(work) + ) + set_seed(config.seed) + device = pick_device(config.device) + labels = load_label_space() + model, tokenizer = load_model_and_tokenizer(config, labels) + train_split = load_split("train") + sampler = start_measurement(device) + times: list[float] = [] + out_dir = work / config.run_name + trainer = Trainer( + model=model, # type: ignore[arg-type] + args=training_arguments(config, out_dir, device), # type: ignore[arg-type] + train_dataset=to_hf_dataset(train_split, tokenizer, config.max_length), # type: ignore[arg-type] + data_collator=DataCollatorWithPadding(tokenizer), # type: ignore[arg-type] + processing_class=tokenizer, # type: ignore[arg-type] + callbacks=[sampler_callback(sampler), step_timer(times)], # type: ignore[list-item] + ) + started = time.perf_counter() + result = trainer.train() + wall = time.perf_counter() - started + losses = [e["loss"] for e in trainer.state.log_history if "loss" in e] + final = out_dir / "final" + trainer.save_model(str(final)) + val = predict_logits(final, load_split("validation"), config) + pred, oos = val.logits.argmax(axis=1), labels.oos_intent_id + steady = times[5:] if len(times) > 10 else times + record = { + "config": config_path, + "model_name": config.model_name, + "model_revision": config.model_revision, + "learning_rate": lr, + "train_batch_size": config.train_batch_size, + "max_length": config.max_length, + "steps": result.global_step, + "wall_seconds": round(wall, 3), + "step_seconds_median": round(float(np.median(steady)), 4), + "step_seconds_mean_after_5": round(float(np.mean(steady)), 4), + "step_seconds_first": round(times[0], 4), + "estimated_full_5_epoch_seconds": round(float(np.median(steady)) * AC2_EPOCH_STEPS, 1), + "loss_first_5": losses[:5], + "loss_last_5": losses[-5:], + "loss_decreased": float(np.mean(losses[-10:])) < float(np.mean(losses[:10])), + "any_nan_loss": any(not math.isfinite(x) for x in losses), + "validation_logits_finite": bool(np.isfinite(val.logits).all()), + "validation_in_scope_accuracy_150": in_scope_accuracy(pred, val.labels, oos), + "validation_oos_recall_151": oos_recall(pred, val.labels, oos), + "parameters": count_parameters(model), + "attn_implementation": getattr(model.config, "_attn_implementation", None), + "peak_memory": memory_report(device, sampler), + "environment": environment(config), + "plan_note": f"full 5 epochs = {plan_steps(replace(config, max_steps=-1), 15_250)}", + } + return record, final + + +def main(argv: list[str] | None = None) -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--config", required=True) + parser.add_argument("--steps", type=int, default=50) + parser.add_argument("--learning-rate", type=float, default=5e-5) + parser.add_argument("--work", default="checkpoints/compat") + parser.add_argument("--out", required=True) + parser.add_argument("--keep-weights", action="store_true", help="for the ONNX export trial") + args = parser.parse_args(argv) + work = Path(args.work) + record, final = run_trial(args.config, args.steps, args.learning_rate, work) + out = Path(args.out) + out.parent.mkdir(parents=True, exist_ok=True) + out.write_text(json.dumps(record, indent=2) + "\n", encoding="utf-8") + print(json.dumps({k: v for k, v in record.items() if k != "environment"}, indent=2)) + if args.keep_weights: + print(f"weights kept at {final}") + else: + shutil.rmtree(final.parent) + print(f"deleted {final.parent}") + + +if __name__ == "__main__": + main() diff --git a/scripts/export_onnx.py b/scripts/export_onnx.py new file mode 100644 index 0000000..f03be8e --- /dev/null +++ b/scripts/export_onnx.py @@ -0,0 +1,155 @@ +"""Export a fine-tuned classifier to ONNX and compare ONNX Runtime (CPU) with PyTorch. + +Compatibility check for RQ7 (docs/PLAN.md section 4): can this encoder be +exported at all, and do the exported logits match? Tries the +``torch.onnx.export`` exporters in order (dynamo, then the TorchScript +one) and records every attempt, including the error message of a failure. +Scores ``--rows`` validation queries with PyTorch (CPU, float32) and with +ONNX Runtime (CPUExecutionProvider) and reports the maximum absolute logit +difference. No quantization here; that is step 6. + +Needs the optional group: ``uv run --group onnx python scripts/export_onnx.py ...``. +""" + +from __future__ import annotations + +import argparse +import json +import platform +import time +import traceback +from pathlib import Path + +import numpy as np + +from tinyrouter.data import load_split + +EXPORTERS = ("dynamo", "torchscript") + + +def encode(tokenizer: object, texts: list[str], max_length: int) -> dict[str, np.ndarray]: + batch = tokenizer( # type: ignore[operator] + texts, truncation=True, max_length=max_length, padding=True, return_tensors="np" + ) + return {"input_ids": batch["input_ids"], "attention_mask": batch["attention_mask"]} + + +def export(model: object, inputs: dict[str, np.ndarray], path: Path, exporter: str) -> None: + import torch + + args = (torch.from_numpy(inputs["input_ids"]), torch.from_numpy(inputs["attention_mask"])) + names = ["input_ids", "attention_mask"] + if exporter == "dynamo": + batch, seq = torch.export.Dim("batch"), torch.export.Dim("sequence") + shapes = {name: {0: batch, 1: seq} for name in names} + torch.onnx.export( + model, args, str(path), input_names=names, output_names=["logits"], + dynamic_shapes=shapes, dynamo=True, + ) # fmt: skip + else: + axes = {name: {0: "batch", 1: "sequence"} for name in names} + axes["logits"] = {0: "batch"} + torch.onnx.export( + model, args, str(path), input_names=names, output_names=["logits"], + dynamic_axes=axes, dynamo=False, opset_version=17, + ) # fmt: skip + + +def onnx_logits(path: Path, inputs: dict[str, np.ndarray]) -> np.ndarray: + import onnxruntime as ort + + session = ort.InferenceSession(str(path), providers=["CPUExecutionProvider"]) + wanted = {i.name for i in session.get_inputs()} + feed = {k: v.astype(np.int64) for k, v in inputs.items() if k in wanted} + return np.asarray(session.run(["logits"], feed)[0]) + + +def torch_logits(model: object, inputs: dict[str, np.ndarray]) -> np.ndarray: + import torch + + with torch.inference_mode(): + return model(**{k: torch.from_numpy(v) for k, v in inputs.items()}).logits.numpy() # type: ignore[operator] + + +def max_diff(a: np.ndarray, b: np.ndarray) -> float: + return float(np.abs(a.astype(np.float64) - b.astype(np.float64)).max()) + + +def attempt( + model: object, inputs: dict, reference: np.ndarray, other: dict, path: Path, exporter: str +) -> dict: + """Export once from ``inputs``; check on them and on ``other`` (another batch and length).""" + started = time.perf_counter() + try: + export(model, inputs, path, exporter) + logits = onnx_logits(path, inputs) + other_diff = max_diff(onnx_logits(path, other), torch_logits(model, other)) + except Exception as exc: # noqa: BLE001 - every failure is the result being recorded + return { + "exporter": exporter, + "ok": False, + "error": f"{type(exc).__name__}: {exc}"[:2000], + "traceback_tail": traceback.format_exc().splitlines()[-8:], + } + weights = sum(p.stat().st_size for p in path.parent.glob("*.data")) + return { + "exporter": exporter, + "ok": True, + "seconds": round(time.perf_counter() - started, 1), + "onnx_bytes": path.stat().st_size + weights, + "max_abs_logit_diff": max_diff(logits, reference), + "other_shape": list(other["input_ids"].shape), + "other_shape_max_abs_logit_diff": other_diff, + "argmax_agree": int((logits.argmax(1) == reference.argmax(1)).sum()), + "rows": int(reference.shape[0]), + } + + +def main(argv: list[str] | None = None) -> None: + import onnx + import onnxruntime + import torch + import transformers + from transformers import AutoModelForSequenceClassification, AutoTokenizer + + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model-dir", required=True) + parser.add_argument("--max-length", type=int, default=64) + parser.add_argument("--rows", type=int, default=10) + parser.add_argument("--out", required=True) + args = parser.parse_args(argv) + model_dir = Path(args.model_dir) + tokenizer = AutoTokenizer.from_pretrained(model_dir) + model = AutoModelForSequenceClassification.from_pretrained(model_dir).eval() + texts = list(load_split("validation").texts[: args.rows]) + inputs = encode(tokenizer, texts, args.max_length) + reference = torch_logits(model, inputs) + longest = max(load_split("validation").texts, key=len) + other = encode(tokenizer, [longest, texts[0], "hi"], args.max_length) + attempts = [] + for exporter in EXPORTERS: + out_dir = model_dir.parent / f"onnx-{exporter}" + out_dir.mkdir(exist_ok=True) + attempts.append(attempt(model, inputs, reference, other, out_dir / "model.onnx", exporter)) + print(json.dumps(attempts[-1], indent=2)) + record = { + "model_type": model.config.model_type, + "rows": len(texts), + "padded_sequence_length": int(inputs["input_ids"].shape[1]), + "attempts": attempts, + "versions": { + "torch": torch.__version__, + "transformers": transformers.__version__, + "onnx": onnx.__version__, + "onnxruntime": onnxruntime.__version__, + }, + "machine": platform.platform(), + } + out = Path(args.out) + out.parent.mkdir(parents=True, exist_ok=True) + out.write_text(json.dumps(record, indent=2) + "\n", encoding="utf-8") + print(f"wrote {out}") + + +if __name__ == "__main__": + main() diff --git a/uv.lock b/uv.lock index 22e5d38..a0bcafd 100644 --- a/uv.lock +++ b/uv.lock @@ -272,6 +272,14 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/29/33/af0635ab07fe83b1788a1dbe370ff3e226062495a998335cb18a1cac81aa/filelock-4.0.1-py3-none-any.whl", hash = "sha256:481a321a27bef441e23c53371c6abc8d7d16e26b97090074ba44f7538a3fd55a", size = 106219, upload-time = "2026-09-19T01:08:14.49Z" }, ] +[[package]] +name = "flatbuffers" +version = "25.12.19" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e8/2d/d2a548598be01649e2d46231d151a6c56d10b964d94043a335ae56ea2d92/flatbuffers-25.12.19-py2.py3-none-any.whl", hash = "sha256:7634f50c427838bb021c2d66a3d1168e9d199b0607e6329399f04846d42e20b4", size = 26661, upload-time = "2025-12-19T23:16:13.622Z" }, +] + [[package]] name = "frozenlist" version = "1.8.0" @@ -531,6 +539,22 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/b3/38/89ba8ad64ae25be8de66a6d463314cf1eb366222074cfda9ee839c56a4b4/mdurl-0.1.2-py3-none-any.whl", hash = "sha256:84008a41e51615a49fc9966191ff91509e3c40b939176e643fd50a5c2196b8f8", size = 9979, upload-time = "2022-08-14T12:40:09.779Z" }, ] +[[package]] +name = "ml-dtypes" +version = "0.6.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/12/72/307d7c4bd0600601c7133fba5cb78af7db968152951c1cd473abb1cda782/ml_dtypes-0.6.0.tar.gz", hash = "sha256:5e60251d32ced5598972e4d5e06a2f044341f9291402551a3f6f0ec44f9299b0", size = 3032327, upload-time = "2026-08-13T14:14:40.215Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/84/6a/441eb053b078954f7fea284dfb288701884d0a1404d39babb858e1649023/ml_dtypes-0.6.0-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:5359c588cc62de6f78d7430f06b65853d884955494d86d6ad90b6dd64a3f3a08", size = 565447, upload-time = "2026-08-13T14:14:01.737Z" }, + { url = "https://files.pythonhosted.org/packages/ed/cf/87e8a6c57eed63a91782a0d229856ddf73e138ce004dd71e2799a9dcdb33/ml_dtypes-0.6.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:37da32aa97749251025666d62372775019594577b9c9e9cfda83bed48d778fdb", size = 360227, upload-time = "2026-08-13T14:14:02.938Z" }, + { url = "https://files.pythonhosted.org/packages/c7/f9/7d76c1eae866f5d4636401b31b6d6dd90e4b4ced1fa7cfdfcca9c60e4bd3/ml_dtypes-0.6.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:3b4a480aa8fd54a1805b8ac10f3f91763926a74f73c0c364c10f9231854f4170", size = 409890, upload-time = "2026-08-13T14:14:04.248Z" }, + { url = "https://files.pythonhosted.org/packages/ba/db/9c61ec2760b5cbfb1c6558d5c991a6d8fd3271053c32db20506a9a90272b/ml_dtypes-0.6.0-cp312-cp312-win_amd64.whl", hash = "sha256:2a3e9d53925597fbffafd2a37048dadeddd0bdaba58058f6ae0869ed709a184d", size = 439333, upload-time = "2026-08-13T14:14:05.501Z" }, + { url = "https://files.pythonhosted.org/packages/6a/57/780ca3e5ab135b9fbdd8e5441abf5f801b30398371b691291e05ab9834c0/ml_dtypes-0.6.0-cp312-cp312-win_arm64.whl", hash = "sha256:6eaed129a4afe90694b8685e2f9b6294849f5eda4af9a15be83a4326eeebd775", size = 552268, upload-time = "2026-08-13T14:14:06.866Z" }, +] + [[package]] name = "mpmath" version = "1.3.0" @@ -621,6 +645,78 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/22/43/1764aff32e4652526ae2f71fa8b3efd8d25c8a3d6926914454e47138ed1e/numpy-2.5.3-cp312-cp312-win_arm64.whl", hash = "sha256:ccb32e0525d29e8b0572eb84c9a57af0e7a4e615726927506f55063c62414034", size = 10485432, upload-time = "2026-09-06T16:24:57.278Z" }, ] +[[package]] +name = "onnx" +version = "1.23.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "ml-dtypes" }, + { name = "numpy" }, + { name = "protobuf" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/29/3a/b73bfb8f506a7b5be43eebf71dbdd351c4116f8100e1103a91bfcbcd3502/onnx-1.23.0.tar.gz", hash = "sha256:3f4b7323ea3397c63aa6ff5d43abe16013ffefc0a96cb8a0be6527ff1e1add40", size = 6022619, upload-time = "2026-09-18T19:30:37.511Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/63/f9/54d8047c2f409cfc314a0df520c6e5703afe78b1395e8c44287724c4b37e/onnx-1.23.0-cp312-abi3-macosx_13_0_universal2.whl", hash = "sha256:6cc93aa89502acbf453a16d79b465d47d9c454f3a3cb3fc5507a8cb59ab53208", size = 9725329, upload-time = "2026-09-18T19:30:07.486Z" }, + { url = "https://files.pythonhosted.org/packages/c6/ff/a7d3fe41debd5b8947182a632b4e001bd3a7a103212e58ef6ca0c3daab91/onnx-1.23.0-cp312-abi3-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:95b77466b4f6bc7c2c63d6004b8db7d963fb17bf24aff13007752e9d2ebf9390", size = 8640232, upload-time = "2026-09-18T19:30:10.058Z" }, + { url = "https://files.pythonhosted.org/packages/09/86/63cdd5740f566fd9de2cadc106529adee0566f2fcafe69a65767eedc1681/onnx-1.23.0-cp312-abi3-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:f336004196a22fbdc16c62e7f26f20635af1826db6147c80ff3b4b8d428fc7ef", size = 8881347, upload-time = "2026-09-18T19:30:12.792Z" }, + { url = "https://files.pythonhosted.org/packages/cd/24/7844929ef08aee3f9cb857d5ce34bb60d44496c9cc5c66cdfd8255917283/onnx-1.23.0-cp312-abi3-pyemscripten_2026_0_wasm32.whl", hash = "sha256:05ad734997f09fcc9633fc2e3715033239792fc9f12aec6ec0c48e32b214e47e", size = 7314556, upload-time = "2026-09-18T19:30:15.456Z" }, + { url = "https://files.pythonhosted.org/packages/52/68/ee1ad44118d0031e28a1cabb1cc02a4c9e30273bdc25571c5847aec5e4e8/onnx-1.23.0-cp312-abi3-win32.whl", hash = "sha256:221037cb9ffe67b701fb180ca871f1ee9089f431d944f3d0313ffae68c356a80", size = 7736116, upload-time = "2026-09-18T19:30:18.06Z" }, + { url = "https://files.pythonhosted.org/packages/81/5a/766abb1411b9f9b06dc4e8fe9ff195d3b715535677c362ab3d950fa0e93f/onnx-1.23.0-cp312-abi3-win_amd64.whl", hash = "sha256:70a2f930b221f9dbdff62704838ce8bf81442787b25048f8b9cdc9118168799e", size = 7872197, upload-time = "2026-09-18T19:30:20.619Z" }, + { url = "https://files.pythonhosted.org/packages/94/34/483fa56b4d47a5210d64dccb522a60b46e0953ca3f72d75b3c4759266f5f/onnx-1.23.0-cp312-abi3-win_arm64.whl", hash = "sha256:dcd4d269818a13c48444c96b7f17c7539d8cda7810ef802431393b6cf5df794b", size = 8037224, upload-time = "2026-09-18T19:30:22.682Z" }, +] + +[[package]] +name = "onnx-ir" +version = "1.0.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "ml-dtypes" }, + { name = "numpy" }, + { name = "onnx" }, + { name = "sympy" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/d6/c2/61194cec0dbc5622273c0ebd592d37cc1dca0d7f1a744f02edd45ac905a3/onnx_ir-1.0.0.tar.gz", hash = "sha256:9e261f25fde8da9612ae5cb43b3b374d5ff469c04af0363cad588b2bb000b812", size = 163121, upload-time = "2026-08-11T14:49:46.895Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/91/cd/6d1637172eb59c7b18ac90ed089d1f599a11fe0e63b4db2d017f3bb38a32/onnx_ir-1.0.0-py3-none-any.whl", hash = "sha256:e578f0d608d3062866b48223616eb2d10a6d6d01f8b8faac596129034f483cc7", size = 185849, upload-time = "2026-08-11T14:49:45.524Z" }, +] + +[[package]] +name = "onnxruntime" +version = "1.30.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "flatbuffers" }, + { name = "numpy" }, + { name = "packaging" }, + { name = "protobuf" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/31/6f/48169f2e62b405bff5053cbd1d73fb5ce41ef7ecd13bb3bfcc191e689b8a/onnxruntime-1.30.0-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:001ed726c9bd5e2bc92faade7d37d889e9606a350b7d5529f0227df2e3bb57fd", size = 21544867, upload-time = "2026-09-10T16:31:19.876Z" }, + { url = "https://files.pythonhosted.org/packages/16/bd/cbc5b8f91963689fdd622f463508c01d0aa95d3f944747b1e0b1eb2160b8/onnxruntime-1.30.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:6c32a000d5139a38ba9349030b0032e3331acb559d596b22738d9d2b343a2b83", size = 21345202, upload-time = "2026-09-10T16:31:23.361Z" }, + { url = "https://files.pythonhosted.org/packages/34/35/e7f862dbacbc99fadd9b14a614e49c99bf0f35fd9927a82f096e3de33531/onnxruntime-1.30.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:fa688e7891a6aa206636fe7372e27ee75fd17713289f6b4fc7b190e0a7de9328", size = 23585654, upload-time = "2026-09-10T16:31:26.65Z" }, + { url = "https://files.pythonhosted.org/packages/a6/13/0f1699f6de549c9324bc9112a2a85b14c517904cd11b562a654643b755a1/onnxruntime-1.30.0-cp312-cp312-win_amd64.whl", hash = "sha256:f3501472571f1b1eee50e017851e7929f5ea37312d2d8c2494a19e8fc58b4a38", size = 14311470, upload-time = "2026-09-10T16:31:31.273Z" }, + { url = "https://files.pythonhosted.org/packages/a2/17/02b13e5f51461f0453b18ab854e2d0bc1b6ec353241b05a2ee6b79e35d87/onnxruntime-1.30.0-cp312-cp312-win_arm64.whl", hash = "sha256:dc4c706f1935ebb62356e6a095b047859badd854482c40560888e95c328ed262", size = 14175072, upload-time = "2026-09-10T16:31:34.541Z" }, +] + +[[package]] +name = "onnxscript" +version = "0.7.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "ml-dtypes" }, + { name = "numpy" }, + { name = "onnx" }, + { name = "onnx-ir" }, + { name = "packaging" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/0a/01/3e3fab8d643ca097ea4aa9e51246643699dfaaa0650589744fe44bc46651/onnxscript-0.7.2.tar.gz", hash = "sha256:2c664f6383d10f332a4d47b2876dcab16dba84909fe703656b19abc281fda165", size = 646719, upload-time = "2026-09-09T17:06:44.567Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b9/3b/06260997cdc41138e58718588a6c87d0eb342bbe0dda8a6aae91d163c384/onnxscript-0.7.2-py3-none-any.whl", hash = "sha256:d0e7121c6a1eefd608058928e111cbdb76709f70d269ff0d07aee493bd1d13c9", size = 754215, upload-time = "2026-09-09T17:06:46.442Z" }, +] + [[package]] name = "packaging" version = "26.3" @@ -687,6 +783,21 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/f5/cd/785c64ed382f3f04201870267b02783f63b4678c2acfddc177a3ebcc2727/propcache-0.5.4-py3-none-any.whl", hash = "sha256:62c60aec739ed00124573cce1178138fd690c7676352d67a37328c1cf51d7468", size = 16338, upload-time = "2026-09-16T00:17:13.106Z" }, ] +[[package]] +name = "protobuf" +version = "7.36.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/d9/89/5b8517baa72f84a67b8a307ba953c91057af618bf40bf676f3c03551f8f0/protobuf-7.36.2.tar.gz", hash = "sha256:497d0463ff3316681da6c0b9e8d06cb465d61abce00b613ab42226175644d1bb", size = 512737, upload-time = "2026-09-17T20:07:59.326Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/32/72/98342feb672507c8f3a69e34b4fa8961f608edba5c1a48a6f47156d92cb5/protobuf-7.36.2-cp310-abi3-macosx_10_9_universal2.whl", hash = "sha256:cbc70b17ee27e28894c7fee8bb04be1abead49e936bc70eb60052531eee2079e", size = 456039, upload-time = "2026-09-17T20:07:51.542Z" }, + { url = "https://files.pythonhosted.org/packages/b6/ea/91fdf7c2b8bbd49cde056f00a9df6773532987e1c00fe2830b895af95c7e/protobuf-7.36.2-cp310-abi3-manylinux2014_aarch64.whl", hash = "sha256:e11e1f0180583a2af89db6a2ecd9e8dc40aa6d2988ca175bfd0e6d12ea72d74e", size = 344219, upload-time = "2026-09-17T20:07:52.914Z" }, + { url = "https://files.pythonhosted.org/packages/17/ab/5fd5f8ece73fad885c5a09aa849b32d70472f954ba3a92d3bb5974ea953b/protobuf-7.36.2-cp310-abi3-manylinux2014_s390x.whl", hash = "sha256:f4fee11ec330d238b34a05c9b675f693c20415d1c5bd7d5320cc2f8a798eb9cf", size = 357223, upload-time = "2026-09-17T20:07:53.985Z" }, + { url = "https://files.pythonhosted.org/packages/db/f3/3996583dd2906297a637af12114deddf7658af6e683fedb83be061983fb5/protobuf-7.36.2-cp310-abi3-manylinux2014_x86_64.whl", hash = "sha256:89f23aa53c24553a2416fd4fd1ec06f74fa42b14b546d8883128813f775bbfd2", size = 343223, upload-time = "2026-09-17T20:07:54.931Z" }, + { url = "https://files.pythonhosted.org/packages/fc/1b/dcc64f358fcb51811b58ae40b3d28f820725f116d86487cc20bd4b130701/protobuf-7.36.2-cp310-abi3-win32.whl", hash = "sha256:912c1221170e16c08d1f086762f563dd61ff83c18b5fa6652952dfaded66f728", size = 442998, upload-time = "2026-09-17T20:07:55.826Z" }, + { url = "https://files.pythonhosted.org/packages/8a/55/b77bda4e5e5f5971fb51b07663694690e9afdb9402136c16a522bd621cad/protobuf-7.36.2-cp310-abi3-win_amd64.whl", hash = "sha256:a300819d441e078a5608c0d3c709796bb548136058fda017ae51d425b44fd353", size = 456514, upload-time = "2026-09-17T20:07:57.188Z" }, + { url = "https://files.pythonhosted.org/packages/e4/04/d52c7016b04b6c5108f26691f9d33ec82a9b65d041f1a9c771137693d618/protobuf-7.36.2-py3-none-any.whl", hash = "sha256:bdb3a345d48db958e6ce1f18e508beb0cc981d64f24088427549c866cd039f1e", size = 179806, upload-time = "2026-09-17T20:07:58.211Z" }, +] + [[package]] name = "psutil" version = "7.2.2" @@ -1043,6 +1154,11 @@ dev = [ llm = [ { name = "anthropic" }, ] +onnx = [ + { name = "onnx" }, + { name = "onnxruntime" }, + { name = "onnxscript" }, +] [package.metadata] requires-dist = [ @@ -1064,6 +1180,11 @@ dev = [ { name = "ruff", specifier = "==0.16.8" }, ] llm = [{ name = "anthropic", specifier = "==1.8.0" }] +onnx = [ + { name = "onnx", specifier = "==1.23.0" }, + { name = "onnxruntime", specifier = "==1.30.0" }, + { name = "onnxscript", specifier = "==0.7.2" }, +] [[package]] name = "tokenizers" From edc2f8c8db14f5f58c48ddb5c346764ec3abf48b Mon Sep 17 00:00:00 2001 From: drewOrc <36374426+drewOrc@users.noreply.github.com> Date: Wed, 23 Sep 2026 14:44:16 +0800 Subject: [PATCH 3/4] docs: freeze the learning-curve hyperparameter protocol PLAN section 4.1 records the S_min and learning-rate pilots, AC2 reuse and the test-isolation rule; acceptance criteria are unchanged. DEVLOG records the ModernBERT trial numbers and the ONNX result. README lists the new make targets and modules. --- DEVLOG.md | 48 ++++++++++++++++++++++++++++++++++++++++++++++++ README.md | 22 ++++++++++++++++++---- docs/PLAN.md | 13 +++++++++++++ 3 files changed, 79 insertions(+), 4 deletions(-) diff --git a/DEVLOG.md b/DEVLOG.md index 0dfd2f8..5f49797 100644 --- a/DEVLOG.md +++ b/DEVLOG.md @@ -4,6 +4,54 @@ --- +## 2026-09-23（晚）：步驟 3 準備，抽樣、步數協定、pilot、執行器與 ModernBERT 相容性試跑 + +### 本次工作 / 執行摘要 +- 超參數協定寫進 `docs/PLAN.md` §4.1：步數 = max(S_min, 5 epoch 步數)，S_min 與各 encoder 的 learning rate 都只用 validation pilot 選，選定值手動填進 `configs/curve.yaml`。 +- 新模組：`sampling.py`（k-shot 抽樣，OOS 用寫死對照表）、`steps.py`（步數規劃，訓練後核對實際步數）、`runs.py`（續跑規則與「等價 run」判斷）、`protocol.py`、`pilots.py`、`curves.py`、`baselines.py`；`metrics.wilson_interval`。 +- logits archive 升到 format 2：加三個 split parquet 的 SHA-256；讀 format 1（AC2）時只在 dataset revision 等於鎖定值才從 `SPLIT_FILES` 補上並標註來源，不符就報錯。 +- AC2 重用：BERT 若選中 lr 5e-5，曲線的 k=100 三個點與 lr pilot 的 5e-5 點都重用 AC2，前提是逐欄設定、訓練列數、步數與 seed 都對得上，且 archive 的 SHA-256 三處一致；任何一項不符就重訓。對照已提交的三個 AC2 結果 JSON 測過，S_min 取 100、200、400 或不設都判定等價。 +- 基準：多數類與 TF-IDF centroid 在每個 (k, seed) 的同一份抽樣上計算，存成與 encoder 相同的 archive 與結果 JSON。 +- ModernBERT 相容性試跑（`scripts/compat_trial.py`、`scripts/export_onnx.py`），結果在 `results/compat/`。只跑 50 步與一次 validation，不是正式 pilot 或曲線。 + +### 核心發現 / 數據 +**ModernBERT 在 M4 MPS 上可以訓練。** 兩個模型同設定：batch 32、max_length 64、50 步、lr 5e-5（試跑用佔位值）、seed 42。 + +| | ModernBERT-base | bert-base-uncased | +|---|---:|---:| +| 參數量 | 149.7M | 109.6M | +| 每步時間（中位數，去掉前 5 步） | 0.234 s | 0.139 s | +| 第一步 | 1.13 s | 0.51 s | +| 峰值記憶體（Metal driver，採樣） | 4.80 GiB | 3.18 GiB | +| 峰值記憶體（存活 tensor，採樣） | 2.30 GiB | 1.66 GiB | +| loss（前 5 步 → 後 5 步） | 5.22 → 4.22 | 5.08 → 5.00 | +| NaN | 無 | 無 | + +- attention 走 `sdpa`（MPS 上沒有 Flash Attention 與 unpadding），沒有報錯或退回警告。 +- 每步時間 ModernBERT 約為 BERT 的 1.69 倍。 +- **全量 5 epoch（2,385 步）預估：** 直接乘是 559 s，但同法估 BERT 得 332 s，而 AC2 實測是 822 到 886 s（每步約 0.36 s），短試跑低估約 2.6 倍（原因未查：可能是長時間執行的降頻、每個 epoch 存檔與每步記憶體採樣）。用 AC2 實測乘上 1.69 倍，**ModernBERT 全量每個 seed 約 24 分鐘**，以這個數字排時程。 +- 粗估整條曲線（S_min 以 400 計，每個 seed 約 5,400 步）：ModernBERT 18 點約 2.7 小時，OOS 消融 3 次約 1.2 小時；BERT 若重用 AC2，剩 15 點約 1 小時。 +- **ONNX 匯出成功，兩種 exporter 都可以。** `torch.onnx.export`（dynamo 與 TorchScript）加 ONNX Runtime 1.30 CPU，10 筆 validation（padding 到 14）最大絕對 logit 誤差 2.1e-5，argmax 10/10 一致；另換一批 3 筆、長度 25 的輸入也只差 1.1e-5，動態 batch 與長度可用。模型檔約 600 MB。TorchScript exporter 有 TracerWarning（masking 相關），實測沒出錯，但 RQ7 以 dynamo 為主。int8 量化還沒試，留給步驟 6。**RQ7 可以維持用 ModernBERT。** +- TF-IDF centroid 若直接把餘弦（0 到 1）當 logit，validation 的溫度擬合在 k ≥ 5 撞到搜尋下界 T = 0.05；改用 CLIP 慣例的 100 倍餘弦，T 落在 2.9 到 7.2 之間。這是 seed 42、43 在 scratch 目錄的檢查，不是正式基準結果。 + +### Blockers / 遇到的問題 +- 磁碟剩約 11 GB。試跑權重與兩個 ONNX 檔（共 1.7 GB）已刪除。 +- 短試跑的每步時間與長時間訓練差很多（見上），時程估計以 AC2 實測為準。 +- pilot 重用 AC2 時，判斷等價要讀 AC2 結果 JSON 的 `config` 與 `training` 欄位；那個檔案裡也有 test 數字，但程式只取這兩欄，validation 數字從 archive 只讀 validation 陣列。 + +### Next +- [ ] `make pilot-lr`（ModernBERT 3 次 + BERT 2 次，約 1.7 小時），把選定值填進 `configs/curve.yaml` +- [ ] `make pilot-steps`（6 次 k=5 短訓練），填 S_min +- [ ] `make curve MODEL=bert`、`make curve MODEL=modernbert`、`make oos-ablation` +- [ ] logits 檔附到 GitHub Release + +### Files / Budget +- 新增：`src/tinyrouter/{sampling,steps,runs,protocol,pilots,curves,baselines}.py`、`configs/{modernbert-base,curve}.yaml`、`scripts/{compat_trial,export_onnx}.py`、`results/compat/*.json` +- 相依：optional group `onnx`（onnx 1.23.0、onnxruntime 1.30.0、onnxscript 0.7.2），CI 不安裝 +- API 花費：US$0 + +--- + ## 2026-09-23：AC2 通過（BERT 流程正確性檢查） ### 本次工作 / 執行摘要 diff --git a/README.md b/README.md index e54f46c..6a304ac 100644 --- a/README.md +++ b/README.md @@ -29,6 +29,11 @@ Other targets: | `make train CONFIG=configs/bert-base.yaml SEED=42` | fine-tune; final weights in `checkpoints//final` | | `make evaluate CONFIG=... SEED=...` | score validation and test, write `results/runs/.json` and the logits archive | | `make ac2` | bert-base-uncased, full data, seeds 42/43/44; writes `results/ac2.json`, exits 1 on FAIL; resumes only runs made with the current config, `FORCE=1` clears and reruns all three, including deleting seed 42's kept weights; after `make clean-checkpoints` only `FORCE=1` brings seed 42's weights back | +| `make pilot-lr` | validation-only pilot: each encoder at k=100, seed 42, lr in {1e-5, 2e-5, 5e-5}; writes `results/pilots/lr.json` and prints the choice (BERT at 5e-5 reuses AC2 seed 42 when it is the same run) | +| `make pilot-steps` | validation-only pilot: both encoders at k=5, seed 42, S_min in {100, 200, 400}; writes `results/pilots/steps.json`. Needs the learning rates in `configs/curve.yaml` | +| `make baselines` | majority-class and TF-IDF centroid baselines on every (k, seed) sample, archived like an encoder run | +| `make curve MODEL=bert\|modernbert` | baselines, then 6 values of k x 3 seeds with the lr and S_min from `configs/curve.yaml`; writes `results/curves/.json`; keeps no weights | +| `make oos-ablation` | ModernBERT, k=100, no OOS training rows, 3 seeds; writes `results/curves/oos-ablation.json` | | `make verify-logits` | check every archive in `results/logits/` against `results/logits-manifest.json` | | `make report` | build `results/summary.md` from `results/runs/*.json` | | `make clean-checkpoints` | delete all trained weights | @@ -44,17 +49,19 @@ Put `HF_TOKEN` and `ANTHROPIC_API_KEY` in a `.env` (copy `.env.example`) if you - **Labels.** The model is trained on all 151 CLINC150 intents (150 plus `oos`). Each intent maps to one of 8 routing targets (7 agents plus `oos`) through `src/tinyrouter/resources/intent_to_agent.json`. - **Data.** `clinc/clinc_oos`, `plus` config, from the Hugging Face Hub at a pinned commit. Each parquet file is checked against its SHA-256 and row count (15,250 / 3,100 / 5,500), and the label names inside it must match the committed `intent_names.json`. - **Calibration.** Temperature scaling, fitted on validation logits only. `fit_temperature` raises `LeakageError` if given anything else, and a test checks that. -- **Logits archive.** Every evaluation stores per-example validation and test logits (float32) with gold labels and metadata (model and revision, seed, k, OOS training rows, dataset revision, label-space hash, git commit, time) in `results/logits/.npz`. Later analysis reads only these files. They are not committed (about 5 MB per bert-base run); their SHA-256 goes into the committed `results/logits-manifest.json`, and the files are attached to a GitHub Release. +- **Logits archive.** Every evaluation stores per-example validation and test logits (float32) with gold labels and metadata (model and revision, seed, k, OOS training rows, dataset revision and the SHA-256 of the three split files, label-space hash, git commit, time) in `results/logits/.npz`. Later analysis reads only these files. They are not committed (about 5 MB per bert-base run); their SHA-256 goes into the committed `results/logits-manifest.json`, and the files are attached to a GitHub Release. - **Training cost.** Each results JSON records parameter counts, training wall time, steps, device and peak memory. How memory is measured depends on the device (sampled Metal driver memory on MPS, the allocator peak on CUDA, process peak RSS on CPU) and is written next to the number; see `src/tinyrouter/efficiency.py`. -- **Metrics.** 8-way accuracy, 150-way in-scope accuracy and OOS recall (as defined in Larson et al. 2019), ECE, and NLL. They are computed before and after calibration on both splits. +- **Learning-curve protocol.** Each curve point trains for max(S_min, 5 epochs) steps. The learning rate of each encoder and the one S_min are chosen by validation-only pilots (`docs/PLAN.md` section 4.1) and written into `configs/curve.yaml` by hand. The pilot code raises `LeakageError` if handed the test split. +- **Metrics.** 8-way accuracy, 150-way in-scope accuracy and OOS recall (as defined in Larson et al. 2019), ECE, and NLL. They are computed before and after calibration on both splits. `wilson_interval` gives the confidence interval for OOS recall. ## Layout ``` -configs/ run configs (bert-base.yaml, smoke.yaml); unknown keys are an error +configs/ run configs (bert-base.yaml, modernbert-base.yaml, smoke.yaml; unknown keys + are an error) and curve.yaml, the pilot-chosen lr and S_min docs/PLAN.md research questions and acceptance criteria docs/OPERATIONS.md CI jobs, branch protection, rollback -scripts/ check_em_dash.py +scripts/ check_em_dash.py; compat_trial.py and export_onnx.py (encoder trial runs) src/tinyrouter/ data.py pinned download, checksum, per-intent subsampling labels.py 151 intents <-> 8 routing targets @@ -64,6 +71,13 @@ src/tinyrouter/ evaluate.py logits, temperature fit, metrics, results JSON archive.py per-example logits archive (.npz) and SHA-256 manifest efficiency.py parameter counts, wall time, peak memory per device + sampling.py k-shot training samples (k per intent, ceil(2.5k) oos from a fixed table) + steps.py training steps = max(S_min, epoch steps), checked after training + runs.py resume rules and when an existing run counts as the same run + protocol.py configs/curve.yaml: pilot-chosen lr per encoder and S_min + pilots.py validation-only lr and S_min pilots + curves.py learning curves and the OOS ablation (reuses AC2 at k=100 when equivalent) + baselines.py majority-class and TF-IDF centroid baselines per curve point ac2.py AC2 run over three seeds and PASS/FAIL verdict llm.py Claude Haiku zero-shot router (baseline and fallback; not run yet) report.py results/*.json -> results/summary.md diff --git a/docs/PLAN.md b/docs/PLAN.md index cad5004..e7541ac 100644 --- a/docs/PLAN.md +++ b/docs/PLAN.md @@ -72,6 +72,19 @@ TinyRouter 問三件事：**小模型要多少標註資料才夠？它知不知 | 推論優化 | encoder 用 Optimum 匯出 ONNX，ONNX Runtime 動態 int8 量化；CPU、batch 1 量延遲 | TensorRT：需要 NVIDIA GPU；只量 PyTorch：沒有優化可講 | | 服務化 | FastAPI `/route` 端點（encoder ONNX + 門檻 + 可關閉的 Haiku fallback）、`/healthz`、Dockerfile；本機壓測 | 只交 notebook：沒有「上線」證據；上雲長期運行：L1 專案不值得維運成本 | +### 4.1 超參數協定（學習曲線，2026-09-23 定案） + +所有 pilot 只看 validation，選出的值凍結後，整條曲線與兩個 encoder 共用。設定檔 `configs/curve.yaml`；pilot 只印出選定值，由人手動填入，程式不自動改設定。 + +| 項目 | 協定 | +|---|---| +| 訓練步數 | 實際步數 = max(S_min, 5 個 epoch 的步數)。S_min 從 {100, 200, 400} 選一個，所有 k 與兩個 encoder 共用。結果 JSON 記錄預定步數、實際步數與由哪一邊決定（`epochs` 或 `min_train_steps`） | +| S_min pilot（`make pilot-steps`） | k=5、seed 42，兩個 encoder 各用自己選定的 learning rate 都跑；選兩者 val 150 類 in-scope 準確率平均最高者，平手取較小的 S_min | +| learning rate pilot（`make pilot-lr`） | 每個 encoder 各自從 {1e-5, 2e-5, 5e-5} 選，k=100、seed 42；標準是 val 150 類 in-scope 準確率，平手看 val OOS recall，再平手取較小的 learning rate。k=100 的 epoch 步數（2,385）大於任何 S_min，所以這個 pilot 不設 S_min，先跑它 | +| AC2 重用 | BERT 的 5e-5 直接重用 AC2 seed 42 的 validation 數字；曲線上 BERT 的 k=100 點，若選中 5e-5，重用 AC2 的三個 run。重用前逐欄比對設定、訓練列數與步數，任何一項不同就重訓 | +| test 隔離 | pilot 不得產生或讀取 test 的任何結果：pilot 的評估函式在程式層級只接受 validation，傳入 test 就 raise，並有測試 | +| 信賴區間 | OOS recall 的信賴區間用 Wilson score interval（`metrics.wilson_interval`，預設 95%） | + **訓練次數估計**：ModernBERT 6 刻度 × 3 seeds + OOS 0 筆消融 × 3 seeds = 21 次；BERT 6 刻度 × 3 seeds = 18 次（其中 k=100 那 3 次就是 AC2）；共 39 次 encoder，k ≤ 25 的每次數十秒；另加 LoRA（全量 + 10-shot）× 3 seeds = 6 次、QLoRA 全量 × 1 seed = 1 次。encoder 小刻度每次數十秒、全量數分鐘到二十分鐘；LoRA 全量預估每次一小時上下，試跑後更新。 ## 5. 驗收條件 From b687d540437d5617fff9056b56dfc6f3d5c7e940 Mon Sep 17 00:00:00 2001 From: drewOrc <36374426+drewOrc@users.noreply.github.com> Date: Wed, 23 Sep 2026 15:05:38 +0800 Subject: [PATCH 4/4] Address review: PR text check, validation-only reuse, sample pinning - pr-text.yml checks the PR title and body (the squash commit's text) on open, edit, push and reopen, through env only; it shares check-disallowed-text.sh and the pattern file with commit-hygiene. - Pilots reuse AC2 without reading test: the archive is hashed whole, only its validation arrays are read, and the record keeps no metrics. - Weights are deleted only after the results JSON, manifest and archive agree; the curve index checks the archive of every trained point. - The curve index compares each run's sample fingerprint with curve_sample now; golden fingerprints pin the sampler in tests. - equivalent_run also requires the same torch and transformers releases; the schedule comparison has its own tests. - TF-IDF scaling documented; RQ3 uses only temperature-scaled signals and the margin for it, and leaves the majority baseline out. - scipy is a direct dev dependency. --- .github/disallowed-trailers.txt | 4 ++ .github/scripts/check-disallowed-text.sh | 23 ++++++ .github/workflows/ci.yml | 16 ++--- .github/workflows/pr-text.yml | 43 +++++++++++ DEVLOG.md | 28 ++++++++ docs/OPERATIONS.md | 3 + docs/PLAN.md | 1 + pyproject.toml | 3 + src/tinyrouter/baselines.py | 37 ++++++---- src/tinyrouter/curves.py | 71 ++++++++++++++++-- src/tinyrouter/pilots.py | 13 ++-- src/tinyrouter/runs.py | 89 +++++++++++++++++++---- tests/run_fakes.py | 20 +++++- tests/test_baselines.py | 21 ++++++ tests/test_curves.py | 92 +++++++++++++++++++++++- tests/test_pilots.py | 35 +++++++++ tests/test_runs.py | 92 +++++++++++++++++++++++- tests/test_sampling.py | 36 ++++++++++ tests/test_text_hygiene.py | 82 +++++++++++++++++++++ uv.lock | 2 + 20 files changed, 660 insertions(+), 51 deletions(-) create mode 100755 .github/scripts/check-disallowed-text.sh create mode 100644 .github/workflows/pr-text.yml create mode 100644 tests/test_text_hygiene.py diff --git a/.github/disallowed-trailers.txt b/.github/disallowed-trailers.txt index c6ac35c..1064123 100644 --- a/.github/disallowed-trailers.txt +++ b/.github/disallowed-trailers.txt @@ -7,3 +7,7 @@ co-authored-by:.*(anthropic|claude|openai|chatgpt|copilot|gemini|cursor) generated (with|by) \[?(claude|chatgpt|copilot|gemini|cursor) noreply@anthropic\.com +# Canary: exercises the checks (for example in a throwaway draft PR) +# without writing a real attribution line that would then stay in +# GitHub's edit history. +tinyrouter-hygiene-canary diff --git a/.github/scripts/check-disallowed-text.sh b/.github/scripts/check-disallowed-text.sh new file mode 100755 index 0000000..001605f --- /dev/null +++ b/.github/scripts/check-disallowed-text.sh @@ -0,0 +1,23 @@ +#!/usr/bin/env bash +# Reads text on stdin and fails if any line matches a pattern in +# .github/disallowed-trailers.txt (extended regex, case-insensitive). +# The one definition shared by the commit-hygiene job (ci.yml) and the +# pr-text-hygiene job (pr-text.yml), so the two cannot drift apart. +# Usage: | check-disallowed-text.sh +# Exit: 0 clean, 1 a line matched, 2 the pattern file has no patterns. +set -euo pipefail + +label="${1:-text}" +patterns_file="${PATTERNS_FILE:-.github/disallowed-trailers.txt}" +patterns="$(grep -vE '^[[:space:]]*(#|$)' "$patterns_file" | paste -sd'|' - || true)" +if [ -z "$patterns" ]; then + echo "::error::$patterns_file has no patterns" + exit 2 +fi +bad="$(grep -iE "$patterns" || true)" +if [ -n "$bad" ]; then + echo "::error::$label must not carry tool-attribution lines:" + printf '%s\n' "$bad" + exit 1 +fi +echo "clean: $label" diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 897747b..993693f 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -87,7 +87,8 @@ jobs: # checkout on pull_request lands on GitHub's synthetic merge commit). # Pushes to main: the pushed range, which is where a squash merge # message edited in the web UI shows up. Patterns live in - # .github/disallowed-trailers.txt. + # .github/disallowed-trailers.txt; the check itself is + # .github/scripts/check-disallowed-text.sh, shared with pr-text.yml. env: EVENT: ${{ github.event_name }} PR_BASE: ${{ github.event.pull_request.base.sha }} @@ -102,17 +103,8 @@ jobs: else range="$PUSH_BEFORE..$PUSH_AFTER" fi - patterns="$(grep -vE '^[[:space:]]*(#|$)' .github/disallowed-trailers.txt | paste -sd'|' -)" - if [ -z "$patterns" ]; then - echo "::error::.github/disallowed-trailers.txt has no patterns" - exit 1 - fi - bad="$(git log "$range" --format='%h %s%n%b' | grep -iE "$patterns" || true)" - if [ -n "$bad" ]; then - echo "::error::commit messages must not carry tool-attribution trailers:" - echo "$bad" - exit 1 - fi + git log "$range" --format='%h %s%n%b' \ + | .github/scripts/check-disallowed-text.sh "commit messages in $range" echo "clean: $(git rev-list --count "$range") commit(s) in $range" network: diff --git a/.github/workflows/pr-text.yml b/.github/workflows/pr-text.yml new file mode 100644 index 0000000..782d0b9 --- /dev/null +++ b/.github/workflows/pr-text.yml @@ -0,0 +1,43 @@ +name: PR text + +# main only takes squash merges, and a squash commit's title and message +# come from the PR title and body, not from the branch's commits. The +# commit-hygiene job in ci.yml checks the branch's commits only, so this +# job checks the PR title and body with the same patterns +# (.github/disallowed-trailers.txt via .github/scripts/check-disallowed-text.sh). +# +# A separate workflow because it runs on `edited`: editing a PR description +# must re-run this check without cancelling or re-running ci.yml's jobs. +# +# Not covered: text typed into the merge dialog by hand. No PR event carries +# it; the push-to-main run of commit-hygiene sees it, after the fact. +on: + pull_request: + types: [opened, edited, synchronize, reopened] + +concurrency: + group: pr-text-${{ github.event.pull_request.number }} + cancel-in-progress: true + +permissions: + contents: read + +jobs: + pr-text-hygiene: + name: pr-text-hygiene + runs-on: ubuntu-latest + timeout-minutes: 5 + steps: + - name: Checkout (for the pattern file and the check script) + uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + + - name: No tool-attribution lines in the PR title or body + # Title and body are attacker-controlled text. They reach the shell + # only as environment variables, never through ${{ }} inside `run`, + # so nothing in them can be executed. + env: + PR_TITLE: ${{ github.event.pull_request.title }} + PR_BODY: ${{ github.event.pull_request.body }} + run: | + printf '%s\n%s\n' "$PR_TITLE" "$PR_BODY" \ + | .github/scripts/check-disallowed-text.sh "PR title and body" diff --git a/DEVLOG.md b/DEVLOG.md index 5f49797..9e5efdf 100644 --- a/DEVLOG.md +++ b/DEVLOG.md @@ -4,6 +4,34 @@ --- +## 2026-09-23（夜）：PR #7 審查修正 + +### 本次工作 / 執行摘要 +- 新增 `pr-text.yml`（job `pr-text-hygiene`）：squash merge 的 commit 取自 PR 標題與內文，原本的 commit-hygiene 看不到。標題與內文只經 env 傳入。兩個 job 共用 `.github/scripts/check-disallowed-text.sh` 與同一份 pattern 檔。 +- pilot 重用 AC2 改為 validation-only：archive 整檔照算 SHA-256，但只讀 metadata 與 validation 陣列；結果 JSON 只取 `run_name`、`config`、`training`、`logits`、`environment`，不含 `metrics`。 +- 權重刪除前先確認結果 JSON、manifest、archive 三處一致；曲線索引對非重用點也驗 archive。 +- 曲線索引逐點比對「訓練時記下的抽樣指紋」與「現在重抽的 `curve_sample(k, seed)`」，numpy 升版造成樣本改變時會報錯；AC2 重用點補上現算的指紋並標註來源。測試裡另釘了 5 組抽樣指紋。 +- 等價判斷新增：torch 與 transformers 版本必須與目前安裝的相同（忽略 `+cpu` 這類 build 標籤）；schedule 層補獨立測試。 +- TF-IDF 的 100 倍：修正 docstring，PLAN §4.1 註明 RQ3 對 TF-IDF 只報溫度校準後的訊號與 margin，多數類排除在 RQ3 外。 +- scipy 列入 dev group。 + +### 核心發現 / 數據 +- 審查實測（seed 42 val，k=100）：TF-IDF 未校準 MSP 的 AURC 在純餘弦為 0.0854、100 倍為 0.0533；entropy 為 0.1418 與 0.0539。倍數確實改變未校準訊號，所以 RQ3 不用它們。 +- 關於 test 洩漏：Drew 在 PR #6（AC2 結果）看過 BERT lr 5e-5 的 test 數字。lr pilot 的選擇規則是機械式的（只比 validation 的正確筆數，再比 validation OOS 正確筆數，再取較小 lr），看過 test 數字不會影響它選出什麼，所以不構成洩漏。 + +### Blockers / 遇到的問題 +- (無) + +### Next +- [ ] PR #7 合併後，由協調者把 `pr-text-hygiene` 加進 main 的 required checks +- [ ] 其餘同上一則 + +### Files / Budget +- `.github/workflows/pr-text.yml`、`.github/scripts/check-disallowed-text.sh`、`tests/test_text_hygiene.py` +- API 花費：US$0 + +--- + ## 2026-09-23（晚）：步驟 3 準備，抽樣、步數協定、pilot、執行器與 ModernBERT 相容性試跑 ### 本次工作 / 執行摘要 diff --git a/docs/OPERATIONS.md b/docs/OPERATIONS.md index dc12d2d..45733d9 100644 --- a/docs/OPERATIONS.md +++ b/docs/OPERATIONS.md @@ -11,9 +11,12 @@ TinyRouter is a research repository with no deployment target: nothing runs as a | `test` | tracked-files guard, `uv sync --locked`, `make lint`, `make test` (offline) | yes | | `commit-hygiene` | rejects tool-attribution trailers in commit messages (patterns in `.github/disallowed-trailers.txt`) | yes | | `network` | `make test-network` and `make smoke` against the Hugging Face Hub | no | +| `pr-text-hygiene` (`pr-text.yml`) | rejects the same patterns in the PR title and body, on open, edit, push and reopen | yes (added to the required checks once it is on `main`) | `network` is not required because it depends on a third-party service. A Hub outage would otherwise block every merge, including ones that do not touch data loading, and a check people learn to override stops meaning anything. It still runs on every PR, and a red `network` on a PR that changes `data.py`, `smoke.py` or a config is a reason not to merge. The weekly run skips the cache, so it is the one that notices if a pinned upstream dataset or model disappears. +`main` only takes squash merges, and a squash commit's title and message come from the PR title and body, not from the branch's commits. That is why `pr-text-hygiene` exists next to `commit-hygiene`; both use `.github/scripts/check-disallowed-text.sh` and the patterns in `.github/disallowed-trailers.txt`. `pr-text-hygiene` is its own workflow so that editing a PR description re-runs it without cancelling `ci.yml`. One gap is left: text typed by hand into the merge dialog is in no PR event, so CI cannot see it before the merge; the `commit-hygiene` run on the push to `main` finds it afterwards, and the fix is then a revert. + `main` is protected: changes go through a pull request, `test` and `commit-hygiene` must pass, and force pushes and branch deletion are refused. ## Rolling back diff --git a/docs/PLAN.md b/docs/PLAN.md index e7541ac..332beee 100644 --- a/docs/PLAN.md +++ b/docs/PLAN.md @@ -84,6 +84,7 @@ TinyRouter 問三件事：**小模型要多少標註資料才夠？它知不知 | AC2 重用 | BERT 的 5e-5 直接重用 AC2 seed 42 的 validation 數字；曲線上 BERT 的 k=100 點，若選中 5e-5，重用 AC2 的三個 run。重用前逐欄比對設定、訓練列數與步數，任何一項不同就重訓 | | test 隔離 | pilot 不得產生或讀取 test 的任何結果：pilot 的評估函式在程式層級只接受 validation，傳入 test 就 raise，並有測試 | | 信賴區間 | OOS recall 的信賴區間用 Wilson score interval（`metrics.wilson_interval`，預設 95%） | +| 基準的不確定性訊號（RQ3） | TF-IDF centroid 的分數是 100 倍餘弦（讓 validation 溫度擬合有內部最佳解）。這個倍數不影響 argmax 與溫度校準後的機率，但會改變未校準的 MSP、entropy 與原始 ECE，所以 TF-IDF 在 RQ3 只報溫度校準後的訊號與分數 margin。多數類對每筆查詢給同一組分數，所有不確定性訊號都是常數，RQ3 不納入 | **訓練次數估計**：ModernBERT 6 刻度 × 3 seeds + OOS 0 筆消融 × 3 seeds = 21 次；BERT 6 刻度 × 3 seeds = 18 次（其中 k=100 那 3 次就是 AC2）；共 39 次 encoder，k ≤ 25 的每次數十秒；另加 LoRA（全量 + 10-shot）× 3 seeds = 6 次、QLoRA 全量 × 1 seed = 1 次。encoder 小刻度每次數十秒、全量數分鐘到二十分鐘；LoRA 全量預估每次一小時上下，試跑後更新。 diff --git a/pyproject.toml b/pyproject.toml index 4c2cb96..c457f2c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -36,6 +36,9 @@ llm = [ dev = [ "pytest==9.1.1", "ruff==0.16.8", + # tests/test_metrics.py checks wilson_interval against scipy. It is also + # a scikit-learn dependency, but a test must not rely on that. + "scipy==1.18.1", ] # ONNX export and ONNX Runtime (RQ7; scripts/export_onnx.py). Not installed # by `uv sync` or in CI; `uv sync --group onnx` installs it. diff --git a/src/tinyrouter/baselines.py b/src/tinyrouter/baselines.py index abb36a9..c1dd2b3 100644 --- a/src/tinyrouter/baselines.py +++ b/src/tinyrouter/baselines.py @@ -15,19 +15,30 @@ frequent agent, ``finance_agent`` (38 of the 150 intents): that is the 8-way majority class, so ``accuracy_8_summed`` is the majority-class accuracy and ``accuracy_8`` (argmax intent, then its agent) always says - oos. Its softmax is the training class distribution, a calibrated - prior; the temperature fit on validation may still move it. -- ``tfidf-centroid``: cosine similarity between the query's TF-IDF vector - and each intent's centroid. The vectorizer (word unigrams and bigrams, - sublinear tf, L2-normalised rows) is fit on the training sample only; - a centroid is the mean of its intent's training vectors, L2-normalised. - The score is ``COSINE_SCALE`` (100) times the cosine, the convention - CLIP uses to turn cosines into logits. The factor changes no argmax; - it only moves the scores to a range where the validation temperature - fit (search range T in [0.05, 100]) has an interior optimum. On the - raw cosines (range [0, 1]) the fitted T hit the lower bound at every - k >= 5 (checked on seed 42 before this was written). A query with no - known term scores 0 everywhere and its argmax is intent 0. + oos. Every query gets the same scores, so every uncertainty signal + (MSP, entropy, margin, with or without temperature) is one constant: + it cannot rank queries, and the RQ3 analysis leaves this baseline out. +- ``tfidf-centroid``: ``COSINE_SCALE`` (100) times the cosine similarity + between the query's TF-IDF vector and each intent's centroid. The + vectorizer (word unigrams and bigrams, sublinear tf, L2-normalised + rows) is fit on the training sample only; a centroid is the mean of its + intent's training vectors, L2-normalised. A query with no known term + scores 0 everywhere and its argmax is intent 0. + + The factor is not neutral. It changes nothing that depends on the + argmax (accuracy, OOS recall) and nothing computed after temperature + scaling or directly on the scores: the validation fit returns a T 100 + times larger, so softmax(scores / T) is the same distribution (a test + checks this), and a margin between two scores only changes by the same + constant factor, so its ranking of queries does not change. It does + change everything computed from softmax(scores) without a temperature: + raw MSP, raw entropy, raw ECE and raw NLL. On seed 42 validation, + raw-softmax AURC at k=100 was 0.0854 (MSP) and 0.1418 (entropy) on + plain cosines and 0.0533 and 0.0539 at 100 times. So for this baseline + RQ3 reports only temperature-scaled signals and the score margin + (docs/PLAN.md section 4.1). The factor was chosen so that the + temperature fit (T in [0.05, 100]) has an interior optimum: on plain + cosines it hit the lower bound at every k >= 5. """ from __future__ import annotations diff --git a/src/tinyrouter/curves.py b/src/tinyrouter/curves.py index 16b4ca0..e828fba 100644 --- a/src/tinyrouter/curves.py +++ b/src/tinyrouter/curves.py @@ -21,6 +21,14 @@ the planned steps and which rule decided them (``epochs`` or ``min_train_steps``), the steps actually run, and ``reused_from`` when the point is an AC2 run. + +Every point's training sample is checked against ``curve_sample(k, seed)`` +recomputed now (``sample_fingerprint``). The encoder's resume check +compares configs only, so a numpy release that changed the sampler's +random stream would otherwise pair old encoder runs with new baseline +samples without notice. AC2 runs predate the fingerprint field; for them +the index records the fingerprint computed now and says so in +``train_sample_sha256_source``. """ from __future__ import annotations @@ -28,6 +36,8 @@ import argparse import json import os +from collections.abc import Callable +from functools import cache from pathlib import Path from tinyrouter.baselines import run_all as run_baselines @@ -38,20 +48,58 @@ EvaluateFn, Log, TrainFn, + archive_intact, default_evaluate, default_train, read_record, reusable_equivalent, run_one, ) +from tinyrouter.sampling import curve_sample, sample_fingerprint from tinyrouter.steps import plan_steps +FingerprintFn = Callable[[RunConfig], str] +BACKFILLED_SAMPLE = "computed at index time from curve_sample(k, seed); the run predates the field" + + +@cache +def _fingerprint(k: int, seed: int, oos_train: int | None) -> str: + return sample_fingerprint(curve_sample(k, seed, oos_train)) + + +def expected_fingerprint(config: RunConfig) -> str: + """Fingerprint of the rows ``config`` should train on, drawn with the numpy installed now.""" + assert config.k_shot is not None + return _fingerprint(config.k_shot, config.seed, config.oos_train) + class CurveError(RuntimeError): """A finished curve point does not look like the run it was supposed to be.""" -def index_entry(config: RunConfig, record: dict, reused_from: str | None) -> dict[str, object]: +def checked_sample( + config: RunConfig, record: dict, reused_from: str | None, fingerprint_fn: FingerprintFn +) -> tuple[str, str]: + """(fingerprint, source): the run's recorded sample must be the one drawn now.""" + expected = fingerprint_fn(config) + recorded = record["training"].get("train_sample_sha256") + if recorded is None and reused_from is not None: + return expected, BACKFILLED_SAMPLE + if recorded != expected: + raise CurveError( + f"{record['run_name']}: trained on sample {recorded}, but curve_sample(" + f"{config.k_shot}, {config.seed}) now gives {expected}; the sampler's output " + "changed (numpy upgrade?), so this run and the baselines no longer share rows" + ) + return expected, "recorded by train()" + + +def index_entry( + config: RunConfig, + record: dict, + reused_from: str | None, + fingerprint_fn: FingerprintFn | None = None, +) -> dict[str, object]: training, logits = record["training"], record["logits"] plan = plan_steps(config, training["train_rows"]) if training["global_step"] != plan.planned_steps: @@ -61,6 +109,10 @@ def index_entry(config: RunConfig, record: dict, reused_from: str | None) -> dic ) if reused_from is None and training.get("step_plan") != plan.as_dict(): raise CurveError(f"{record['run_name']}: recorded step plan differs from this config's") + if reused_from is None and not archive_intact(RunPaths.of(config), record, config.seed): + raise CurveError(f"{record['run_name']}: archive does not match its record and manifest") + fingerprint_fn = fingerprint_fn or expected_fingerprint + fingerprint, source = checked_sample(config, record, reused_from, fingerprint_fn) return { "k": config.k_shot, "seed": config.seed, @@ -70,7 +122,8 @@ def index_entry(config: RunConfig, record: dict, reused_from: str | None) -> dic "logits_sha256": logits["sha256"], "train_rows": training["train_rows"], "oos_train_rows": training["oos_train_rows"], - "train_sample_sha256": training.get("train_sample_sha256"), + "train_sample_sha256": fingerprint, + "train_sample_sha256_source": source, "planned_steps": plan.planned_steps, "epoch_steps": plan.epoch_steps, "decided_by": plan.decided_by, @@ -80,20 +133,25 @@ def index_entry(config: RunConfig, record: dict, reused_from: str | None) -> dic def run_point( - config: RunConfig, protocol: CurveProtocol, train_fn: TrainFn, evaluate_fn: EvaluateFn, log: Log + config: RunConfig, + protocol: CurveProtocol, + train_fn: TrainFn, + evaluate_fn: EvaluateFn, + log: Log, + fingerprint_fn: FingerprintFn | None, ) -> dict[str, object]: donor = ac2_donor(protocol, config) if donor is not None: record = reusable_equivalent(config, donor) if record is not None: log(f"{config.run_name}: reusing {donor.run_name} (equivalent run, archive intact)") - return index_entry(config, record, donor.run_name) + return index_entry(config, record, donor.run_name, fingerprint_fn) log(f"{config.run_name}: {donor.run_name} is not an equivalent intact run; training") run_one(config, train_fn, evaluate_fn, log, keep_weights=False) record = read_record(RunPaths.of(config)) if record is None: raise CurveError(f"{config.run_name}: no results JSON after evaluation") - return index_entry(config, record, None) + return index_entry(config, record, None, fingerprint_fn) def run_curve( @@ -103,9 +161,10 @@ def run_curve( train_fn: TrainFn = default_train, evaluate_fn: EvaluateFn = default_evaluate, log: Log = print, + fingerprint_fn: FingerprintFn | None = None, ) -> dict[str, object]: """Run or resume every config, then write results/curves/.json.""" - entries = [run_point(c, protocol, train_fn, evaluate_fn, log) for c in configs] + entries = [run_point(c, protocol, train_fn, evaluate_fn, log, fingerprint_fn) for c in configs] shas = [e["logits_sha256"] for e in entries] if len(set(shas)) != len(shas): raise CurveError(f"{name}: two points share one logits archive") diff --git a/src/tinyrouter/pilots.py b/src/tinyrouter/pilots.py index 5fdb205..c76cd5e 100644 --- a/src/tinyrouter/pilots.py +++ b/src/tinyrouter/pilots.py @@ -3,8 +3,10 @@ ``make pilot-lr``: each encoder at k=100, seed 42, learning rate in {1e-5, 2e-5, 5e-5}. Chosen by validation 150-way in-scope accuracy, then validation OOS recall, then the smaller learning rate. BERT at 5e-5 is the -AC2 seed-42 run, reused when ``runs.equivalent_run`` confirms it; only its -validation logits are read. +AC2 seed-42 run, reused when ``runs.equivalent_run`` confirms it. From +that run only the validation arrays and metadata of its archive are read +(the whole file is hashed), and only the results-JSON fields in +``runs.VALIDATION_ONLY_FIELDS``, which exclude ``metrics``. ``make pilot-steps``: both encoders at k=5, seed 42, with their chosen learning rates, S_min in {100, 200, 400}. Chosen by the mean of the two @@ -12,7 +14,8 @@ Neither pilot scores, loads or writes anything from the test split: ``predict_validation`` and ``validation_scores`` raise ``LeakageError`` on -any other split, the reused AC2 archive is read through +any other split, the reused AC2 archive and record are read through +``reusable_equivalent(..., validation_only=True)`` and ``load_validation_logits``, and nothing here calls ``evaluate()`` (which scores test). Outputs go to ``results/pilots/{lr,steps}.json``. The selected values are printed, not written into ``configs/curve.yaml``; @@ -172,7 +175,9 @@ def point_entry(point: PilotPoint, training: dict, scores: dict, source: str) -> def reused_entry(point: PilotPoint, protocol: CurveProtocol) -> dict | None: """The AC2 run as this point, if it is an equivalent run; validation logits only.""" donor = ac2_donor(protocol, point.config) - record = None if donor is None else reusable_equivalent(point.config, donor) + record = ( + None if donor is None else reusable_equivalent(point.config, donor, validation_only=True) + ) if donor is None or record is None: return None _, val = load_validation_logits(RunPaths.of(donor).logits) diff --git a/src/tinyrouter/runs.py b/src/tinyrouter/runs.py index e8a6288..673686b 100644 --- a/src/tinyrouter/runs.py +++ b/src/tinyrouter/runs.py @@ -10,6 +10,11 @@ under a different config: would training with ``config`` have run the same code on the same rows with the same hyperparameters? See its docstring for the two rewrites it accepts; everything else must be equal. + +The pilots use ``reusable_equivalent(..., validation_only=True)``: it +hashes the whole archive file (a hash is not a result), reads only the +validation arrays and metadata from it, and keeps only the record fields +in ``VALIDATION_ONLY_FIELDS``, none of which holds a test number. """ from __future__ import annotations @@ -18,9 +23,16 @@ import shutil from collections.abc import Callable from dataclasses import fields +from importlib.metadata import version from pathlib import Path -from tinyrouter.archive import ArchiveError, load_logits, read_manifest, remove_from_manifest +from tinyrouter.archive import ( + ArchiveError, + load_logits, + load_validation_logits, + read_manifest, + remove_from_manifest, +) from tinyrouter.config import ADDED_FIELD_DEFAULTS, LOCATION_FIELDS, RunConfig from tinyrouter.data import sha256_of from tinyrouter.evaluate import RunPaths @@ -36,14 +48,28 @@ DATA_AND_SCHEDULE_FIELDS = frozenset( {"per_intent", "k_shot", "oos_train", "num_train_epochs", "max_steps", "min_train_steps"} ) +# Record fields the validation-only path keeps; ``metrics`` (which has test) is dropped. +VALIDATION_ONLY_FIELDS = ("run_name", "config", "training", "logits", "environment") +# Libraries whose version must match for a recorded run to count as this code's run. +PINNED_LIBRARIES = ("torch", "transformers") + + +class RunIncompleteError(RuntimeError): + """Evaluation returned but the run's results JSON and archive do not check out.""" def run_dir(config: RunConfig) -> Path: return Path(config.checkpoint_root) / config.run_name -def archive_intact(paths: RunPaths, record: dict[str, object], seed: int) -> bool: - """The archive's SHA-256 agrees in results JSON, manifest and file, and it is this seed's.""" +def archive_intact( + paths: RunPaths, record: dict[str, object], seed: int, validation_only: bool = False +) -> bool: + """The archive's SHA-256 agrees in results JSON, manifest and file, and it is this seed's. + + ``validation_only`` checks the seed through the metadata and validation + arrays alone, without reading the test arrays. + """ if not paths.logits.exists(): return False logits = record.get("logits") @@ -52,17 +78,23 @@ def archive_intact(paths: RunPaths, record: dict[str, object], seed: int) -> boo if not recorded or not recorded == listed == sha256_of(paths.logits): return False try: - archive = load_logits(paths.logits) + if validation_only: + metadata, _ = load_validation_logits(paths.logits) + else: + metadata = load_logits(paths.logits).metadata except ArchiveError: return False - return archive.metadata["seed"] == seed + return metadata["seed"] == seed -def read_record(paths: RunPaths) -> dict[str, object] | None: +def read_record(paths: RunPaths, keep: tuple[str, ...] | None = None) -> dict[str, object] | None: + """The results JSON, or only the ``keep`` fields of it.""" if not paths.results_json.exists(): return None record = json.loads(paths.results_json.read_text(encoding="utf-8")) - return record if isinstance(record, dict) else None + if not isinstance(record, dict): + return None + return record if keep is None else {k: record[k] for k in keep if k in record} def is_done(config: RunConfig) -> bool: @@ -113,6 +145,11 @@ def run_one( log(f"{config.run_name}: training") final_dir = train_fn(config) evaluate_fn(config, final_dir) + if not is_done(config): + raise RunIncompleteError( + f"{config.run_name}: evaluation returned but the results JSON, manifest and " + f"archive do not agree; weights kept in {run_dir(config)} for a rerun" + ) log(f"{config.run_name}: evaluated, logits archived") if not keep_weights and run_dir(config).exists(): shutil.rmtree(run_dir(config)) @@ -183,12 +220,15 @@ def equivalent_run(config: RunConfig, record: dict[str, object]) -> bool: All other fields must be equal. The record's own training summary must also show the rows and steps this config plans: train rows, oos rows, - and global_step. + and global_step; and the record must have been made with the torch and + transformers versions installed now (``same_libraries``). """ recorded = config_from_record(record.get("config")) training = record.get("training") if recorded is None or not isinstance(training, dict): return False + if not same_libraries(record.get("environment")): + return False mine, theirs = training_identity(config), training_identity(recorded) if mine is None or theirs is None or mine != theirs: return False @@ -202,6 +242,24 @@ def equivalent_run(config: RunConfig, record: dict[str, object]) -> bool: ) +def public_version(text: str) -> str: + """Drop a local build label: the Linux CPU wheel ``2.14.0+cpu`` is release 2.14.0.""" + return text.split("+", 1)[0] + + +def same_libraries(environment: object) -> bool: + """The recorded torch and transformers releases are the ones installed now.""" + if not isinstance(environment, dict): + return False + for name in PINNED_LIBRARIES: + recorded = environment.get(name) + if not isinstance(recorded, str): + return False + if public_version(recorded) != public_version(version(name)): + return False + return True + + def default_train(config: RunConfig) -> Path: from tinyrouter.train import prepare_train_split, train @@ -214,12 +272,19 @@ def default_evaluate(config: RunConfig, model_dir: Path) -> dict[str, object]: return evaluate(config, model_dir) -def reusable_equivalent(config: RunConfig, donor: RunConfig) -> dict[str, object] | None: - """The donor run's record if it is done, intact, and equivalent to ``config``; else None.""" +def reusable_equivalent( + config: RunConfig, donor: RunConfig, validation_only: bool = False +) -> dict[str, object] | None: + """The donor run's record if it is done, intact, and equivalent to ``config``; else None. + + With ``validation_only`` the returned record holds only + ``VALIDATION_ONLY_FIELDS`` and the archive's test arrays are never read. + """ paths = RunPaths.of(donor) - record = read_record(paths) + record = read_record(paths, VALIDATION_ONLY_FIELDS if validation_only else None) if record is None or record.get("run_name") != donor.run_name: return None if not equivalent_run(config, record): return None - return record if archive_intact(paths, record, config.seed) else None + intact = archive_intact(paths, record, config.seed, validation_only) + return record if intact else None diff --git a/tests/run_fakes.py b/tests/run_fakes.py index f370891..a711f7f 100644 --- a/tests/run_fakes.py +++ b/tests/run_fakes.py @@ -18,24 +18,39 @@ from tinyrouter.steps import plan_steps +def fake_fingerprint(config: RunConfig) -> str: + """What FakeRuns records as the sample fingerprint; tests hand it to the curve index.""" + return f"sample-{config.k_shot}-{config.seed}-{config.oos_train}" + + +def current_environment() -> dict[str, str]: + from importlib.metadata import version + + return {"torch": version("torch"), "transformers": version("transformers")} + + def summary_for(config: RunConfig) -> dict[str, object]: if config.k_shot is None: rows, oos = 15_250, 250 else: rows, oos = planned_rows(config.k_shot, config.oos_train) plan = plan_steps(config, rows) - return { + summary = { "run_name": config.run_name, "seed": config.seed, "config": asdict(config), "train_rows": rows, "oos_train_rows": oos, "k_shot": config.k_shot, - "train_sample_sha256": f"sample-{config.k_shot}-{config.seed}", + "train_sample_sha256": fake_fingerprint(config), "step_plan": plan.as_dict(), "global_step": plan.planned_steps, "train_wall_seconds": 1.0, } + if config.k_shot is None: + # Like the AC2 summaries, written before the fingerprint field existed. + del summary["train_sample_sha256"] + return summary class FakeRuns: @@ -71,6 +86,7 @@ def evaluate(self, config: RunConfig, model_dir: Path) -> dict[str, object]: record = { "run_name": config.run_name, "config": asdict(config), + "environment": current_environment(), "training": training, "logits": {"file": paths.logits.name, "sha256": entry["sha256"]}, "metrics": score(splits["validation"], splits["test"]), diff --git a/tests/test_baselines.py b/tests/test_baselines.py index 8034984..3240b2f 100644 --- a/tests/test_baselines.py +++ b/tests/test_baselines.py @@ -127,3 +127,24 @@ def test_baselines_see_exactly_the_rows_the_encoder_of_that_point_trains_on(tmp_ def test_unknown_baseline_is_refused(): with pytest.raises(ValueError, match="unknown baseline"): baselines.baseline_logits("knn", TRAIN, EVALS) + + +def test_temperature_scaling_undoes_the_cosine_scale(): + """1x and 100x cosines, each with its own validation-fitted T, give the same probabilities.""" + from tinyrouter.calibrate import SplitLogits, apply_temperature, fit_temperature + + rng = np.random.default_rng(0) + n = 400 + gold = rng.integers(0, 151, size=n) + cosine = rng.uniform(0.0, 0.35, size=(n, 151)) + cosine[np.arange(n), gold] += rng.uniform(0.0, 0.1, size=n) + plain = SplitLogits("validation", cosine, gold) + scaled = SplitLogits("validation", baselines.COSINE_SCALE * cosine, gold) + t_plain, t_scaled = fit_temperature(plain), fit_temperature(scaled) + assert 0.05 < t_plain < 1 and t_scaled == pytest.approx(100 * t_plain, rel=1e-4) + np.testing.assert_allclose( + apply_temperature(plain.logits, t_plain), + apply_temperature(scaled.logits, t_scaled), + rtol=1e-4, + atol=1e-8, + ) diff --git a/tests/test_curves.py b/tests/test_curves.py index 45f28b9..26fcf18 100644 --- a/tests/test_curves.py +++ b/tests/test_curves.py @@ -4,7 +4,7 @@ import pytest -from run_fakes import FakeRuns +from run_fakes import FakeRuns, fake_fingerprint from test_pilots import make_protocol from tinyrouter.curves import CurveError, index_entry, run_curve from tinyrouter.evaluate import RunPaths @@ -14,6 +14,14 @@ pytestmark = pytest.mark.filterwarnings("ignore::tinyrouter.calibrate.TemperatureBoundWarning") +@pytest.fixture(autouse=True) +def fake_samples(monkeypatch): + """The fake runs record fake_fingerprint; the index compares against the same function.""" + import tinyrouter.curves as curves + + monkeypatch.setattr(curves, "expected_fingerprint", fake_fingerprint) + + def quiet(_: str) -> None: pass @@ -192,3 +200,85 @@ def test_the_ablation_cli_skips_the_baselines(tmp_path, monkeypatch): curves.main(["--ablation"]) assert calls == ["curve oos-ablation"] assert Path(tmp_path / "res" / "curves" / "oos-ablation.json").exists() + + +def test_index_refuses_a_non_reused_point_whose_archive_changed(tmp_path): + protocol = make_protocol(tmp_path, s_min=400, bert_lr=5e-5, modern_lr=2e-5) + config = protocol.curve_config("modernbert", 5, 42) + fake = FakeRuns() + run_one(config, fake.train, fake.evaluate, quiet, keep_weights=False) + record = json.loads(RunPaths.of(config).results_json.read_text()) + assert index_entry(config, record, None)["run_name"] == config.run_name + archive = RunPaths.of(config).logits + archive.write_bytes(archive.read_bytes() + b"x") + with pytest.raises(CurveError, match="archive does not match"): + index_entry(config, record, None) + + +def test_curve_writes_no_index_when_an_evaluation_leaves_a_broken_archive(tmp_path): + from tinyrouter.runs import RunIncompleteError + + protocol = make_protocol(tmp_path, s_min=400, bert_lr=5e-5, modern_lr=2e-5) + fake = FakeRuns() + + def truncating(config, model_dir): + record = fake.evaluate(config, model_dir) + archive = RunPaths.of(config).logits + archive.write_bytes(archive.read_bytes()[:100]) + return record + + with pytest.raises(RunIncompleteError): + run_curve( + "modernbert", + protocol.curve_configs("modernbert"), + protocol, + fake.train, + truncating, + quiet, + ) + assert not (tmp_path / "res" / "curves" / "modernbert.json").exists() + assert list((tmp_path / "ckpt").rglob("model.safetensors")) + + +def test_a_point_trained_on_another_sample_than_curve_sample_now_gives_is_refused(tmp_path): + protocol = make_protocol(tmp_path, s_min=400, bert_lr=5e-5, modern_lr=2e-5) + configs = protocol.curve_configs("modernbert")[:3] + fake = FakeRuns() + run_curve("modernbert", configs, protocol, fake.train, fake.evaluate, quiet) + + def drifted(config): + return fake_fingerprint(config) + "-numpy-upgrade" + + with pytest.raises(CurveError, match="sampler's output changed"): + run_curve("modernbert", configs, protocol, fake.train, fake.evaluate, quiet, drifted) + + +def test_reused_ac2_points_get_the_fingerprint_computed_now_and_say_so(tmp_path): + from tinyrouter.curves import BACKFILLED_SAMPLE + + protocol = make_protocol(tmp_path, s_min=400, bert_lr=5e-5, modern_lr=2e-5) + ac2_runs(protocol) + configs = [c for c in protocol.curve_configs("bert") if c.k_shot == 100] + fake = FakeRuns() + body = run_curve("bert", configs, protocol, fake.train, fake.evaluate, quiet) + for point, config in zip(body["points"], configs, strict=True): + assert point["reused_from"] is not None + assert point["train_sample_sha256"] == fake_fingerprint(config) + assert point["train_sample_sha256_source"] == BACKFILLED_SAMPLE + + +def test_the_real_fingerprint_function_draws_the_curve_sample(monkeypatch): + import tinyrouter.curves as curves + import tinyrouter.sampling as sampling + from test_sampling import FULL + from tinyrouter.sampling import sample_fingerprint, sample_k_shot + + monkeypatch.setattr(sampling, "load_split", lambda name: FULL) + curves._fingerprint.cache_clear() + try: + assert curves._fingerprint(5, 43, None) == sample_fingerprint(sample_k_shot(FULL, 5, 43)) + assert curves._fingerprint(100, 42, 0) == sample_fingerprint( + sample_k_shot(FULL, 100, 42, oos_override=0) + ) + finally: + curves._fingerprint.cache_clear() diff --git a/tests/test_pilots.py b/tests/test_pilots.py index dc4ca08..19b8428 100644 --- a/tests/test_pilots.py +++ b/tests/test_pilots.py @@ -216,3 +216,38 @@ def broken(config, model_dir) -> SplitLogits: again = FakeRuns() pilots.run_pilot("lr", protocol, tmp_path / "lr.json", again.train, again.validation, quiet) assert len(again.trained) == len(pilots.lr_points(protocol)) - 1 + + +def test_pilot_reads_no_test_array_from_the_ac2_archive(tmp_path, monkeypatch): + from test_runs import spy_on_npz_reads + + protocol = make_protocol(tmp_path) + ac2_donor_run(protocol, in_scope_correct=2_902) + seen = spy_on_npz_reads(monkeypatch) + fake = FakeRuns() + body = pilots.run_pilot( + "lr", protocol, tmp_path / "lr.json", fake.train, fake.validation, quiet + ) + assert any(p["source"].startswith("reused") for p in body["points"]) + assert seen and not [key for key in seen if key.startswith("test")] + + +def test_pilot_runs_and_chooses_the_same_without_the_donors_metrics(tmp_path): + protocol = make_protocol(tmp_path) + ac2_donor_run(protocol, in_scope_correct=2_990) + fake = FakeRuns() + fake.correct = {(MODERN_NAME, 2e-5): 2_950} + first = pilots.run_pilot( + "lr", protocol, tmp_path / "a.json", fake.train, fake.validation, quiet + ) + donor_json = tmp_path / "res" / "runs" / "bert-base-uncased-full-seed42.json" + record = json.loads(donor_json.read_text()) + del record["metrics"] + donor_json.write_text(json.dumps(record)) + again = FakeRuns() + again.correct = fake.correct + second = pilots.run_pilot( + "lr", protocol, tmp_path / "b.json", again.train, again.validation, quiet + ) + assert first["selected"] == second["selected"] == {"bert": 5e-5, "modernbert": 2e-5} + assert len(again.trained) == 5 diff --git a/tests/test_runs.py b/tests/test_runs.py index a9871c1..5e08d26 100644 --- a/tests/test_runs.py +++ b/tests/test_runs.py @@ -6,7 +6,7 @@ import pytest -from run_fakes import FakeRuns +from run_fakes import FakeRuns, current_environment from tinyrouter.config import ADDED_FIELD_DEFAULTS, load_config from tinyrouter.evaluate import RunPaths from tinyrouter.runs import equivalent_run, reusable_equivalent, run_one @@ -32,6 +32,7 @@ def ac2_like_record(seed: int = 42, training: dict | None = None) -> dict: return { "run_name": BERT.with_seed(seed).run_name, "config": config, + "environment": current_environment(), "training": { "seed": seed, "train_rows": 15_250, @@ -108,6 +109,7 @@ def test_min_train_steps_that_decides_equals_the_same_fixed_max_steps(): record = { "run_name": fixed.run_name, "config": asdict(fixed), + "environment": current_environment(), "training": {"seed": 42, "train_rows": 763, "oos_train_rows": 13, "global_step": 200}, } assert equivalent_run(by_floor, record) @@ -140,3 +142,91 @@ def test_reusable_equivalent_needs_an_intact_archive(tmp_bert): def test_reusable_equivalent_is_none_when_the_donor_never_ran(tmp_bert): point = replace(tmp_bert, k_shot=100, min_train_steps=400) assert reusable_equivalent(point, tmp_bert) is None + + +@pytest.mark.parametrize( + "fields", + [{"max_steps": 2385}, {"num_train_epochs": 4.999}], + ids=["fixed_2385_steps", "epochs_4.999"], +) +def test_same_step_count_under_another_schedule_is_not_the_same_run(fields): + """Both give 2,385 steps on the full split, but through other TrainingArguments.""" + record = ac2_like_record() + record["config"].update(fields) + assert record["training"]["global_step"] == 2385 + assert not equivalent_run(curve_point(), record) + + +@pytest.mark.parametrize("library", ["torch", "transformers"]) +def test_a_run_recorded_with_another_library_release_is_not_the_same_run(library): + record = ac2_like_record() + record["environment"][library] = "0.0.1" + assert not equivalent_run(curve_point(), record) + + +def test_a_record_without_library_versions_is_not_the_same_run(): + record = ac2_like_record() + del record["environment"] + assert not equivalent_run(curve_point(), record) + + +def test_a_local_build_label_does_not_count_as_another_release(): + from tinyrouter.runs import public_version + + record = ac2_like_record() + record["environment"]["torch"] = public_version(record["environment"]["torch"]) + "+cpu" + assert equivalent_run(curve_point(), record) + + +def test_evaluation_that_leaves_a_broken_archive_keeps_the_weights(tmp_bert): + from tinyrouter.runs import RunIncompleteError, run_dir + + fake = FakeRuns() + + def truncating(config, model_dir): + record = fake.evaluate(config, model_dir) + archive = RunPaths.of(config).logits + archive.write_bytes(archive.read_bytes()[:100]) + return record + + point = replace(tmp_bert, k_shot=5, min_train_steps=200) + with pytest.raises(RunIncompleteError, match="weights kept"): + run_one(point, fake.train, truncating, lambda _: None, keep_weights=False) + assert (run_dir(point) / "final" / "model.safetensors").exists() + + +def spy_on_npz_reads(monkeypatch) -> list[str]: + import numpy as np + + seen: list[str] = [] + original = np.lib.npyio.NpzFile.__getitem__ + + def spy(self, key): + seen.append(key) + return original(self, key) + + monkeypatch.setattr(np.lib.npyio.NpzFile, "__getitem__", spy) + return seen + + +def test_validation_only_reuse_reads_no_test_array_and_no_metrics(tmp_bert, monkeypatch): + donor = tmp_bert.with_seed(42) + run_one(donor, FakeRuns().train, FakeRuns().evaluate, lambda _: None, keep_weights=False) + point = replace(tmp_bert, k_shot=100, min_train_steps=None) + seen = spy_on_npz_reads(monkeypatch) + record = reusable_equivalent(point, donor, validation_only=True) + assert record is not None + assert "metrics" not in record + assert seen and not [key for key in seen if key.startswith("test")] + seen.clear() + assert reusable_equivalent(point, donor) is not None + assert "test_logits" in seen # the full check does read test, which is why pilots avoid it + + +def test_validation_only_reuse_still_rejects_a_changed_archive(tmp_bert): + donor = tmp_bert.with_seed(42) + run_one(donor, FakeRuns().train, FakeRuns().evaluate, lambda _: None, keep_weights=False) + archive = RunPaths.of(donor).logits + archive.write_bytes(archive.read_bytes() + b"x") + point = replace(tmp_bert, k_shot=100) + assert reusable_equivalent(point, donor, validation_only=True) is None diff --git a/tests/test_sampling.py b/tests/test_sampling.py index 188e0e1..0d12eff 100644 --- a/tests/test_sampling.py +++ b/tests/test_sampling.py @@ -124,3 +124,39 @@ def test_on_the_real_train_split_k_100_is_everything_and_k_5_is_balanced(): assert sample_k_shot(real, 100, seed=42).texts == real.texts counts = Counter(sample_k_shot(real, 5, seed=42).intents.tolist()) assert counts.pop(OOS) == 13 and set(counts.values()) == {5} + + +def rng_free_train() -> Split: + """OOS+-shaped split built without any random generator, so only the sampler draws.""" + in_scope = [i for i in range(151) if i != OOS] + intents = [i for r in range(100) for i in in_scope] + for n in range(250): + intents.insert(n * 61, OOS) + texts = tuple(f"q{n}-{i}" for n, i in enumerate(intents)) + return Split("train", texts, np.array(intents, dtype=np.int64)) + + +# SHA-256 of sample_fingerprint on rng_free_train(), numpy 2.5.3 (uv.lock). If +# these change, numpy's Generator stream changed: every k < 100 sample differs +# from the one the recorded runs trained on. Do not update the values to make +# the test pass; the curves must be rerun (curves.py also refuses to index such +# runs, by comparing each run's recorded fingerprint with curve_sample now). +GOLDEN_FINGERPRINTS = { + (1, 42): "73b3054b10915c25150fbaf12f841a8dbf9b7d501050100c9b1c1b8ffed9b9eb", + (5, 42): "bcb7915540c4392dbb999b056e7e0e90deeeaa07260b45850b36447f8549cfc1", + (5, 43): "8616df6a42cf0d60893c7b32f64cbd9bcfc99d645c57df25c150442c89e86bfd", + (25, 44): "ebd8fa4ebcaa263c378315dff1a627bf6d49dd53215a0d6d3b4a159266dca1cb", + (50, 42): "7829d6e72d2af9c9a602787ccf1890c19a9d6ddacaff98762c9eaadde178a5b1", +} + + +@pytest.mark.parametrize(("k", "seed"), sorted(GOLDEN_FINGERPRINTS)) +def test_samples_are_pinned_across_numpy_releases(k, seed): + sample = sample_k_shot(rng_free_train(), k, seed) + assert sample_fingerprint(sample) == GOLDEN_FINGERPRINTS[(k, seed)] + + +def test_the_rng_free_split_is_oos_plus_shaped(): + split = rng_free_train() + counts = Counter(split.intents.tolist()) + assert counts.pop(OOS) == 250 and set(counts.values()) == {100} diff --git a/tests/test_text_hygiene.py b/tests/test_text_hygiene.py new file mode 100644 index 0000000..c1e9752 --- /dev/null +++ b/tests/test_text_hygiene.py @@ -0,0 +1,82 @@ +"""The shared attribution check used by both the commit-hygiene and pr-text-hygiene CI jobs.""" + +import re +import subprocess +from pathlib import Path + +import pytest +import yaml + +ROOT = Path(__file__).parent.parent +SCRIPT = ROOT / ".github" / "scripts" / "check-disallowed-text.sh" +WORKFLOWS = ROOT / ".github" / "workflows" + + +def check(text: str, patterns_file: Path | None = None) -> subprocess.CompletedProcess: + env = {"PATH": "/usr/bin:/bin"} + if patterns_file is not None: + env["PATTERNS_FILE"] = str(patterns_file) + return subprocess.run( + ["bash", str(SCRIPT), "test input"], + input=text, + capture_output=True, + text=True, + cwd=ROOT, + env=env, + ) + + +@pytest.mark.parametrize( + "line", + [ + "Co-Authored-By: Hygiene Canary ", + "TINYROUTER-HYGIENE-CANARY", + "co-authored-by: Some Bot ", + ], +) +def test_a_disallowed_line_fails_and_is_printed(line): + result = check(f"Fix the sampler\n\nBody text\n{line}\n") + assert result.returncode == 1 + assert line in result.stdout + + +def test_clean_text_and_human_co_authors_pass(): + result = check("Add pilots\n\nCo-Authored-By: Jane Doe \n") + assert result.returncode == 0, result.stdout + assert "clean" in result.stdout + + +def test_an_empty_pattern_file_is_an_error_not_a_pass(tmp_path): + empty = tmp_path / "patterns.txt" + empty.write_text("# only comments\n\n") + assert check("anything\n", empty).returncode == 2 + + +def steps_of(workflow: str, job: str) -> list[dict]: + return yaml.safe_load((WORKFLOWS / workflow).read_text())["jobs"][job]["steps"] + + +def test_pr_text_job_runs_on_edits_and_passes_title_and_body_through_env_only(): + workflow = yaml.safe_load((WORKFLOWS / "pr-text.yml").read_text()) + # PyYAML reads the bare key `on` as True. + trigger = workflow[True]["pull_request"]["types"] + assert set(trigger) == {"opened", "edited", "synchronize", "reopened"} + step = steps_of("pr-text.yml", "pr-text-hygiene")[-1] + assert step["env"]["PR_TITLE"] == "${{ github.event.pull_request.title }}" + assert step["env"]["PR_BODY"] == "${{ github.event.pull_request.body }}" + assert "${{" not in step["run"] + assert "check-disallowed-text.sh" in step["run"] + + +def test_no_workflow_interpolates_pr_or_commit_text_inside_run(): + risky = re.compile(r"\$\{\{[^}]*(title|body|head_ref|message)[^}]*\}\}") + for path in WORKFLOWS.glob("*.yml"): + for job in yaml.safe_load(path.read_text())["jobs"].values(): + for step in job["steps"]: + assert not risky.search(step.get("run", "")), (path.name, step.get("name")) + + +def test_commit_hygiene_uses_the_same_check(): + run = steps_of("ci.yml", "commit-hygiene")[-1]["run"] + assert "check-disallowed-text.sh" in run + assert "disallowed-trailers.txt" not in run diff --git a/uv.lock b/uv.lock index a0bcafd..8136178 100644 --- a/uv.lock +++ b/uv.lock @@ -1150,6 +1150,7 @@ dependencies = [ dev = [ { name = "pytest" }, { name = "ruff" }, + { name = "scipy" }, ] llm = [ { name = "anthropic" }, @@ -1178,6 +1179,7 @@ requires-dist = [ dev = [ { name = "pytest", specifier = "==9.1.1" }, { name = "ruff", specifier = "==0.16.8" }, + { name = "scipy", specifier = "==1.18.1" }, ] llm = [{ name = "anthropic", specifier = "==1.8.0" }] onnx = [