| """ |
| Construct provenance-aware paired examples for the acceptance experiment. |
| |
| This module intentionally has no MLX import. Pair construction, null matching, |
| and holdout validation can therefore be tested without a Metal device. |
| """ |
|
|
| import hashlib |
| import json |
| from collections import Counter |
|
|
| import numpy as np |
|
|
|
|
| def bootstrap_mean_ci(values, n_boot=2000, seed=0): |
| """Percentile interval for a mean, resampling the supplied cluster units.""" |
| values = np.asarray(values, dtype=np.float64) |
| if values.ndim != 1: |
| raise ValueError("bootstrap values must be one-dimensional") |
| if values.size == 0: |
| return (float("nan"), float("nan")) |
| if not np.all(np.isfinite(values)): |
| raise ValueError("bootstrap values must be finite") |
| rng = np.random.default_rng(seed) |
| idx = rng.integers(0, values.size, size=(n_boot, values.size)) |
| means = values[idx].mean(axis=1) |
| return float(np.percentile(means, 2.5)), float(np.percentile(means, 97.5)) |
|
|
|
|
| def paired_ratio_ci(num_by_doc, den_by_doc, n_boot=4000, seed=0): |
| """ |
| Bootstrap a ratio of arm means while preserving document-level pairing. |
| |
| Each input element is one document's mean acceptance for an arm. Token |
| positions inside a document are correlated and are not resampled as if they |
| were independent. |
| """ |
| num = np.asarray(num_by_doc, dtype=np.float64) |
| den = np.asarray(den_by_doc, dtype=np.float64) |
| if num.ndim != 1 or den.ndim != 1: |
| raise ValueError("paired bootstrap inputs must be one-dimensional") |
| if num.size != den.size: |
| raise ValueError( |
| f"paired bootstrap length mismatch: {num.size} versus {den.size}" |
| ) |
| if num.size == 0: |
| return (float("nan"), float("nan"), float("nan")) |
| if not np.all(np.isfinite(num)) or not np.all(np.isfinite(den)): |
| raise ValueError("paired bootstrap inputs must be finite") |
| den_mean = den.mean() |
| point = float(num.mean() / den_mean) if den_mean else float("nan") |
|
|
| rng = np.random.default_rng(seed) |
| idx = rng.integers(0, num.size, size=(n_boot, num.size)) |
| num_means = num[idx].mean(axis=1) |
| den_means = den[idx].mean(axis=1) |
| with np.errstate(divide="ignore", invalid="ignore"): |
| ratios = num_means / den_means |
| ratios = ratios[np.isfinite(ratios)] |
| if ratios.size == 0: |
| return (point, float("nan"), float("nan")) |
| return ( |
| point, |
| float(np.percentile(ratios, 2.5)), |
| float(np.percentile(ratios, 97.5)), |
| ) |
|
|
|
|
| def paired_ratio_difference_ci( |
| trained_num_by_doc, |
| trained_den_by_doc, |
| control_num_by_doc, |
| control_den_by_doc, |
| n_boot=4000, |
| seed=0, |
| ): |
| """Bootstrap trained minus control ratios over the same target documents.""" |
| arrays = [ |
| np.asarray(values, dtype=np.float64) |
| for values in ( |
| trained_num_by_doc, |
| trained_den_by_doc, |
| control_num_by_doc, |
| control_den_by_doc, |
| ) |
| ] |
| if any(values.ndim != 1 for values in arrays): |
| raise ValueError("control comparison inputs must be one-dimensional") |
| sizes = {values.size for values in arrays} |
| if len(sizes) != 1: |
| raise ValueError("control comparison inputs must have equal length") |
| size = arrays[0].size |
| if size < 2: |
| raise ValueError("control comparison requires at least two documents") |
| if any(not np.all(np.isfinite(values)) for values in arrays): |
| raise ValueError("control comparison inputs must be finite") |
| trained_num, trained_den, control_num, control_den = arrays |
| if trained_den.mean() <= 0 or control_den.mean() <= 0: |
| raise ValueError("control comparison denominators must have positive means") |
|
|
| trained_ratio = float(trained_num.mean() / trained_den.mean()) |
| control_ratio = float(control_num.mean() / control_den.mean()) |
| difference = trained_ratio - control_ratio |
| rng = np.random.default_rng(seed) |
| idx = rng.integers(0, size, size=(n_boot, size)) |
| trained_boot = trained_num[idx].mean(axis=1) / trained_den[idx].mean(axis=1) |
| control_boot = control_num[idx].mean(axis=1) / control_den[idx].mean(axis=1) |
| differences = trained_boot - control_boot |
| differences = differences[np.isfinite(differences)] |
| if differences.size != n_boot: |
| raise ValueError("control comparison bootstrap produced non-finite values") |
| return ( |
| trained_ratio, |
| control_ratio, |
| difference, |
| float(np.percentile(differences, 2.5)), |
| float(np.percentile(differences, 97.5)), |
| ) |
|
|
|
|
| def iter_holdout(path): |
| """Read provenance-bearing holdout records from JSONL.""" |
| document_ids = set() |
| content_hashes = set() |
| with open(path, encoding="utf-8") as f: |
| for line_no, line in enumerate(f, 1): |
| if not line.strip(): |
| continue |
| try: |
| row = json.loads(line) |
| except json.JSONDecodeError as exc: |
| raise ValueError(f"{path}:{line_no}: invalid JSON: {exc}") from exc |
| missing = [ |
| key for key in ("text", "language", "repo", "document_id") |
| if not row.get(key) |
| ] |
| if missing: |
| raise ValueError( |
| f"{path}:{line_no}: missing required fields {missing}" |
| ) |
| digest = hashlib.sha256(row["text"].encode("utf-8")).hexdigest() |
| declared = row.get("content_sha256") |
| if declared and declared != digest: |
| raise ValueError( |
| f"{path}:{line_no}: content_sha256 does not match text" |
| ) |
| if row["document_id"] in document_ids: |
| raise ValueError( |
| f"{path}:{line_no}: duplicate document_id " |
| f"{row['document_id']}" |
| ) |
| if digest in content_hashes: |
| raise ValueError( |
| f"{path}:{line_no}: duplicate document content" |
| ) |
| document_ids.add(row["document_id"]) |
| content_hashes.add(digest) |
| row["content_sha256"] = digest |
| yield row |
|
|
|
|
| def iter_local_records(texts): |
| """Label legacy local text input so reports cannot mistake it for a holdout.""" |
| for i, text in enumerate(texts): |
| yield { |
| "text": text, |
| "language": "unknown", |
| "repo": "local-untracked", |
| "document_id": f"local-{i}", |
| } |
|
|
|
|
| def file_sha256(path): |
| h = hashlib.sha256() |
| with open(path, "rb") as f: |
| for block in iter(lambda: f.read(1024 * 1024), b""): |
| h.update(block) |
| return h.hexdigest() |
|
|
|
|
| def validate_publication_inputs( |
| receipt_path, |
| holdout_path, |
| tokenizer_path, |
| settings, |
| checkpoint_meta, |
| role, |
| ): |
| """Fail closed unless files, settings, and checkpoint match the contract.""" |
| with open(receipt_path, encoding="utf-8") as f: |
| receipt = json.load(f) |
| if receipt.get("schema_version") != 2: |
| raise ValueError("publication receipt must use schema_version 2") |
| contract = receipt.get("analysis_contract") |
| if not isinstance(contract, dict): |
| raise ValueError("receipt has no analysis_contract") |
| if role not in ("trained", "untrained-control"): |
| raise ValueError(f"invalid publication role {role!r}") |
|
|
| problems = [] |
| holdout_hash = file_sha256(holdout_path) |
| tokenizer_hash = file_sha256(tokenizer_path) |
| expected_holdout = receipt.get("clean_holdout", {}).get("sha256") |
| expected_tokenizer = receipt.get("pair_readiness", {}).get( |
| "tokenizer_sha256" |
| ) |
| if holdout_hash != expected_holdout: |
| problems.append( |
| f"holdout sha256 {holdout_hash} != registered {expected_holdout}" |
| ) |
| if tokenizer_hash != expected_tokenizer: |
| problems.append( |
| f"tokenizer sha256 {tokenizer_hash} != registered {expected_tokenizer}" |
| ) |
|
|
| expected_settings = contract.get("settings", {}) |
| for key, expected in expected_settings.items(): |
| actual = settings.get(key) |
| if actual != expected: |
| problems.append(f"{key} {actual!r} != registered {expected!r}") |
| for key in ( |
| "instrument_version", |
| "primary_comparison", |
| "primary_depth", |
| "bootstrap_unit", |
| "sensitivity_subset", |
| ): |
| actual = settings.get(key) |
| expected = contract.get(key) |
| if actual != expected: |
| problems.append(f"{key} {actual!r} != registered {expected!r}") |
|
|
| expected_args = contract.get("model_args", {}) |
| actual_args = checkpoint_meta.get("model_args", {}) |
| for key, expected in expected_args.items(): |
| actual = actual_args.get(key) |
| if actual != expected: |
| problems.append( |
| f"model_args.{key} {actual!r} != registered {expected!r}" |
| ) |
|
|
| checkpoint_cfg = checkpoint_meta.get("config", {}) |
| if checkpoint_cfg.get("seed") != contract.get("model_seed"): |
| problems.append( |
| f"model seed {checkpoint_cfg.get('seed')!r} != registered " |
| f"{contract.get('model_seed')!r}" |
| ) |
| if role == "trained": |
| expected_step = contract.get("trained_checkpoint_step") |
| if checkpoint_meta.get("step") != expected_step: |
| problems.append( |
| f"trained checkpoint step {checkpoint_meta.get('step')!r} " |
| f"!= registered {expected_step!r}" |
| ) |
| expected_run = contract.get("trained_run_name") |
| if checkpoint_cfg.get("run_name") != expected_run: |
| problems.append( |
| f"run name {checkpoint_cfg.get('run_name')!r} " |
| f"!= registered {expected_run!r}" |
| ) |
| else: |
| control = contract.get("untrained_control", {}) |
| if checkpoint_meta.get("step") != control.get("step"): |
| problems.append( |
| f"untrained control step {checkpoint_meta.get('step')!r} " |
| f"!= registered {control.get('step')!r}" |
| ) |
| if checkpoint_cfg.get("lr") != control.get("learning_rate"): |
| problems.append( |
| f"untrained control lr {checkpoint_cfg.get('lr')!r} " |
| f"!= registered {control.get('learning_rate')!r}" |
| ) |
| if checkpoint_cfg.get("run_name") != control.get("run_name"): |
| problems.append( |
| f"untrained control run name " |
| f"{checkpoint_cfg.get('run_name')!r} " |
| f"!= registered {control.get('run_name')!r}" |
| ) |
| if checkpoint_cfg.get("initialization_only") is not control.get( |
| "initialization_only" |
| ): |
| problems.append( |
| f"initialization_only " |
| f"{checkpoint_cfg.get('initialization_only')!r} " |
| f"!= registered {control.get('initialization_only')!r}" |
| ) |
| if checkpoint_meta.get("optimizer_state_included") is not control.get( |
| "optimizer_state_included" |
| ): |
| problems.append( |
| f"optimizer_state_included " |
| f"{checkpoint_meta.get('optimizer_state_included')!r} " |
| f"!= registered {control.get('optimizer_state_included')!r}" |
| ) |
|
|
| if problems: |
| raise ValueError( |
| "publication contract mismatch:\n- " + "\n- ".join(problems) |
| ) |
| return { |
| "receipt_path": receipt_path, |
| "receipt_sha256": file_sha256(receipt_path), |
| "role": role, |
| "holdout_sha256": holdout_hash, |
| "tokenizer_sha256": tokenizer_hash, |
| "registered_at": receipt.get("analysis_contract_registered_at"), |
| }, receipt |
|
|
|
|
| def guard_frozen_holdout_exploration( |
| receipt_path, holdout_path, post_training_acknowledged |
| ): |
| """Prevent an accidental blind-holdout peek during model development.""" |
| with open(receipt_path, encoding="utf-8") as f: |
| receipt = json.load(f) |
| frozen_hash = receipt.get("clean_holdout", {}).get("sha256") |
| supplied_hash = file_sha256(holdout_path) |
| if supplied_hash == frozen_hash and not post_training_acknowledged: |
| raise ValueError( |
| "refusing to expose the frozen publication holdout to an " |
| "exploratory checkpoint before training is frozen; after all " |
| "training decisions are final, add " |
| "--post-training-frozen-holdout" |
| ) |
|
|
|
|
| def validate_pair_summary(receipt, summary): |
| """Check the constructed target and decoy set against its frozen receipt.""" |
| expected = receipt.get("pair_readiness", {}) |
| fields = ( |
| "paired_examples", |
| "repositories", |
| "decoy_match", |
| "language_counts", |
| "unique_decoy_documents", |
| "max_decoy_reuse", |
| "decoy_reuse_histogram", |
| "disjoint_sensitivity_pairs", |
| ) |
| problems = [] |
| for key in fields: |
| if summary.get(key) != expected.get(key): |
| problems.append( |
| f"{key} {summary.get(key)!r} != registered {expected.get(key)!r}" |
| ) |
| if problems: |
| raise ValueError( |
| "constructed pair set does not match receipt:\n- " |
| + "\n- ".join(problems) |
| ) |
|
|
|
|
| def disjoint_pair_indices(metadata): |
| """ |
| Greedily retain pairs whose target and decoy documents have not appeared. |
| |
| Input order is frozen by the registered holdout seed. The resulting |
| sensitivity subset contains no document in more than one target-decoy pair. |
| """ |
| used = set() |
| selected = [] |
| for index, row in enumerate(metadata): |
| target = row["document_id"] |
| decoy = row["decoy_document_id"] |
| if target in used or decoy in used: |
| continue |
| selected.append(index) |
| used.update((target, decoy)) |
| return selected |
|
|
|
|
| def summarize_pair_dependencies(metadata): |
| """Report shuffled-decoy reuse and the document-disjoint subset size.""" |
| reuse = Counter(row["decoy_document_id"] for row in metadata) |
| histogram = Counter(reuse.values()) |
| return { |
| "unique_decoy_documents": len(reuse), |
| "max_decoy_reuse": max(reuse.values(), default=0), |
| "decoy_reuse_histogram": { |
| str(count): documents |
| for count, documents in sorted(histogram.items()) |
| }, |
| "disjoint_sensitivity_pairs": len(disjoint_pair_indices(metadata)), |
| } |
|
|
|
|
| def lexical_profile(tok, ids): |
| """ |
| Small format profile used only to choose a null suffix. |
| |
| Suffix token length is already exact. These character ratios keep the decoy |
| close in layout and lexical texture so the primary comparison is less able |
| to win merely because the null suffix looks like a different kind of file. |
| """ |
| text = tok.decode(ids) |
| n = max(len(text), 1) |
| return np.asarray([ |
| text.count("\n") / n, |
| sum(c.isspace() for c in text) / n, |
| sum(c.isalnum() or c == "_" for c in text) / n, |
| sum(c.isdigit() for c in text) / n, |
| sum(c in "{}[]();,:.=+-*/" for c in text) / n, |
| len(set(ids)) / max(len(ids), 1), |
| ], dtype=np.float64) |
|
|
|
|
| def choose_decoy(staged, index): |
| """ |
| Match a decoy by language, repository independence, and lexical profile. |
| |
| The filters are relaxed only when the available pool cannot satisfy them. |
| Every relaxation is recorded in pair metadata and summarized in the report. |
| """ |
| target = staged[index] |
| others = [(j, item) for j, item in enumerate(staged) if j != index] |
| same_language = [ |
| (j, item) for j, item in others |
| if item["language"] == target["language"] |
| ] |
| language_pool = same_language or others |
| different_repo = [ |
| (j, item) for j, item in language_pool |
| if item["repo"] != target["repo"] |
| ] |
| pool = different_repo or language_pool |
| if not pool: |
| raise ValueError("a shuffled-suffix control requires two documents") |
|
|
| _, decoy = min( |
| pool, |
| key=lambda candidate: ( |
| float(np.linalg.norm( |
| target["suffix_profile"] - candidate[1]["suffix_profile"] |
| )), |
| candidate[0], |
| ), |
| ) |
| if decoy["language"] != target["language"]: |
| quality = "different_language" |
| elif decoy["repo"] == target["repo"]: |
| quality = "same_repo" |
| else: |
| quality = "matched" |
| distance = float(np.linalg.norm( |
| target["suffix_profile"] - decoy["suffix_profile"] |
| )) |
| return decoy, quality, distance |
|
|
|
|
| def build_pairs(records, tok, sentinels, n_examples, prefix_len, span_len, |
| suffix_len, rng): |
| """ |
| One document yields one L2R, FIM, and shuffled-suffix triple. |
| |
| The decoy has the exact same token length as the true suffix. It is selected |
| from the same language, from a different repository when possible, and by |
| nearest lexical profile. A positive claim still requires both FIM over the |
| shuffled null and FIM over L2R. |
| """ |
| pairs = [] |
| need = prefix_len + span_len + suffix_len + 8 |
| staged = [] |
| pool_limit = max(n_examples + 1, n_examples * 4) |
| for record in records: |
| if len(staged) >= pool_limit: |
| break |
| ids = tok.encode(record["text"]).ids |
| if len(ids) < need: |
| continue |
| p0 = int(rng.integers(0, len(ids) - need + 1)) |
| a = p0 + prefix_len |
| b = a + span_len |
| c = b + suffix_len |
| suffix = ids[b:c] |
| staged.append({ |
| "prefix": ids[p0:a], |
| "middle": ids[a:b], |
| "suffix": suffix, |
| "suffix_profile": lexical_profile(tok, suffix), |
| "language": str(record.get("language") or "unknown"), |
| "repo": str(record.get("repo") or "unknown"), |
| "document_id": str(record.get("document_id") or len(staged)), |
| "content_sha256": record.get("content_sha256"), |
| "token_offset": p0, |
| }) |
|
|
| if len(staged) < 2: |
| return [] |
|
|
| for i, target in enumerate(staged[:n_examples]): |
| decoy_item, match_quality, match_distance = choose_decoy(staged, i) |
| prefix = target["prefix"] |
| middle = target["middle"] |
| suffix = target["suffix"] |
| decoy = decoy_item["suffix"] |
|
|
| l2r = prefix + middle |
| fim = ([sentinels["prefix"]] + prefix |
| + [sentinels["suffix"]] + suffix |
| + [sentinels["middle"]] + middle) |
| shuf = ([sentinels["prefix"]] + prefix |
| + [sentinels["suffix"]] + decoy |
| + [sentinels["middle"]] + middle) |
|
|
| pairs.append({ |
| "l2r": (l2r, (len(prefix), len(prefix) + len(middle))), |
| "fim": (fim, (len(fim) - len(middle), len(fim))), |
| "fim_shuf": (shuf, (len(shuf) - len(middle), len(shuf))), |
| "_meta": { |
| "language": target["language"], |
| "repo": target["repo"], |
| "document_id": target["document_id"], |
| "content_sha256": target["content_sha256"], |
| "token_offset": target["token_offset"], |
| "decoy_language": decoy_item["language"], |
| "decoy_repo": decoy_item["repo"], |
| "decoy_document_id": decoy_item["document_id"], |
| "decoy_match": match_quality, |
| "decoy_lexical_distance": match_distance, |
| }, |
| }) |
| return pairs |
|
|
|
|
| def bin_acceptance_by_baseline_nll(acceptance_by_depth, baseline_nlls, edges): |
| """ |
| Bin every treatment arm by the same corresponding L2R token difficulty. |
| |
| At draft depth j+1, acceptance position t targets the token whose baseline |
| NLL is at t+j+1. Using an arm's own NLL would condition on a post-treatment |
| variable because suffix visibility changes that distribution. |
| """ |
| rows = [] |
| for j, per_document in enumerate(acceptance_by_depth): |
| bins = [] |
| for k in range(len(edges) - 1): |
| values = [] |
| for acceptance, baseline_nll in zip( |
| per_document, baseline_nlls): |
| aligned = np.asarray(baseline_nll)[j + 1:] |
| acceptance = np.asarray(acceptance) |
| n = min(len(acceptance), len(aligned)) |
| if n <= 0: |
| continue |
| keep = ( |
| (aligned[:n] >= edges[k]) |
| & (aligned[:n] < edges[k + 1]) |
| ) |
| values.append(acceptance[:n][keep]) |
| selected = np.concatenate(values) if values else np.zeros(0) |
| bins.append({ |
| "bin": k, |
| "nll_range": [ |
| float(edges[k]) if np.isfinite(edges[k]) else None, |
| ( |
| float(edges[k + 1]) |
| if np.isfinite(edges[k + 1]) |
| else None |
| ), |
| ], |
| "acceptance": ( |
| float(selected.mean()) if selected.size else None |
| ), |
| "n": int(selected.size), |
| }) |
| rows.append({"depth": j + 1, "bins": bins}) |
| return rows |
|
|
