diff --git a/README.md b/README.md index 289c2bb..b7aee9c 100644 --- a/README.md +++ b/README.md @@ -142,6 +142,7 @@ uv run python -m cli.generate_task_rollouts \ --temperature 0.7 \ --top_p 0.95 \ --max_new_tokens 1024 \ + --max_wall_time_minutes 120 \ --seed 1000 \ --device mps ``` @@ -162,7 +163,7 @@ uv run python -m cli.generate_task_rollouts \ --device mps ``` -Generation is resumable and records the exact chat template, token IDs, decoding configuration, random seed, code state, model/tokenizer revisions, and aligned generation-time confidence trace. Length-capped reference responses fail validation and must be regenerated. +Generation is resumable and records the exact chat template, token IDs, decoding configuration, random seed, code state, model/tokenizer revisions, and aligned generation-time confidence trace. Each completed rollout is flushed and synced to disk. `--max_wall_time_minutes` stops only at one of those durable boundaries; rerun the identical command without `--overwrite` to continue. Length-capped reference responses fail validation and must be regenerated. `cot_distortion` is a test-only transfer family imported from the complete official MonitorBench evaluated-artifact matrix; it is not generated by `generate_task_rollouts`. Run the pinned MonitorBench source on the evaluated model, copy and fill the immutable run-manifest template, then import its `.tested.jsonl` outputs: @@ -283,6 +284,7 @@ uv run python -m cli.extract_text_embeddings \ --embedding_model_key qwen3_embedding_0_6b \ --views prompt_text,answer_text,transcript_text,response_prefix_text_p10,response_prefix_text_p25,response_prefix_text_p50,response_prefix_text_p75,response_prefix_text_p100 \ --output_dir outputs/text_embeddings/Qwen3-4B/sycophancy \ + --max_wall_time_minutes 120 \ --device mps uv run python -m cli.extract_text_embeddings \ @@ -295,7 +297,7 @@ uv run python -m cli.extract_text_embeddings \ --device mps ``` -Repeat for every source and transfer task. The embedding runner reuses each immutable dataset/view cache across all few-shot seeds. +Repeat for every source and transfer task. The embedding runner reuses each immutable dataset/view cache across all few-shot seeds. Existing compatible view caches are validated and skipped, so the same command can safely continue in later sessions. ## Run the experiments @@ -407,6 +409,7 @@ uv run python -m cli.extract_task_activations \ --trajectory_prefix_percentiles 10,25,50,75,100 \ --trajectory_prefix_stack_views \ --output_dir outputs/final_features/Qwen3-4B/honesty_control \ + --max_wall_time_minutes 120 \ --device mps ``` @@ -541,7 +544,7 @@ The final gate requires at least three monitored-model families, ten seeds, `k={ ## Apple Silicon -All model-backed commands accept `--device auto|cpu|cuda|cuda:N|mps`. On an M-series Mac: +All model-backed commands accept `--device auto|cpu|cuda|cuda:N|mps`. Rollout generation, activation extraction, and text embedding extraction can be divided into bounded sessions with `--max_wall_time_minutes`. Activation extraction writes an atomic checkpoint every 32 examples, and LLM-judge scoring every 16 examples, by default. Repeating the exact command resumes those checkpoints, while `--overwrite` deliberately starts the applicable stage over. On an M-series Mac: ```bash uv run python -m cli.run_frontier_experiments \ diff --git a/cli/build_invariant_shift_scenarios.py b/cli/build_invariant_shift_scenarios.py index 62b7a59..0236179 100644 --- a/cli/build_invariant_shift_scenarios.py +++ b/cli/build_invariant_shift_scenarios.py @@ -96,6 +96,7 @@ def build_invariant_variants( axis: str, registry: dict[str, Any], registry_sha256: str, + eligible_parent_splits: set[str] | None = None, ) -> list[dict[str, Any]]: if axis not in GENERATORS: raise ValueError(f"axis must be one of {sorted(GENERATORS)}") @@ -105,6 +106,8 @@ def build_invariant_variants( transformed_groups: set[str] = set() for parent in parsed: + if eligible_parent_splits and parent["protocol_split"] not in eligible_parent_splits: + continue parent_metadata = parent.get("metadata") or {} parent_falsification = validate_falsification_metadata( parent_metadata.get("falsification"), @@ -244,6 +247,7 @@ def main() -> None: axis=args.axis, registry=registry, registry_sha256=registry_sha256, + eligible_parent_splits={"test"}, ) _write_jsonl(Path(args.output), rows) print(f"saved {len(rows)} base and deterministic shifted scenarios to {args.output}") diff --git a/cli/build_on_policy_scenarios.py b/cli/build_on_policy_scenarios.py index 0731341..aa0ca1d 100644 --- a/cli/build_on_policy_scenarios.py +++ b/cli/build_on_policy_scenarios.py @@ -14,6 +14,7 @@ ) from data.rollout_schema import ScenarioRecord from data.source_adapters import ( + choice_list, ensure_text, first_text, format_mcq_prompt, @@ -103,6 +104,28 @@ def _enforce_heldout_group_splits(rows: list[dict[str, Any]]) -> list[dict[str, return output +def _reserve_overlapping_heldout_groups( + rows: list[dict[str, Any]], split_seed: int +) -> list[dict[str, Any]]: + roles_by_group: dict[str, set[str]] = {} + for row in rows: + template = row["metadata"]["falsification"]["axes"]["template"] + roles_by_group.setdefault(str(row["group_id"]), set()).add(template["role"]) + overlapping = { + group for group, roles in roles_by_group.items() if roles == {"source", "heldout"} + } + return [ + row + for row in rows + if not ( + str(row["group_id"]) in overlapping + and row["metadata"]["falsification"]["axes"]["template"]["role"] + == "heldout" + and _hashed_split(str(row["group_id"]), split_seed) != "test" + ) + ] + + def _honesty_context(row: dict[str, Any]) -> str | None: return first_text( row, @@ -176,6 +199,30 @@ def _honesty_group_prefix(row: dict[str, Any], split_name: str) -> str: return f"mask_{split_name}" +def _sycophancy_targets(base: dict[str, Any]) -> tuple[str, str, str | None, str | None] | None: + choices = choice_list(base) + correct_label = ensure_text(base.get("correct_letter")) + pressure_label = ensure_text(base.get("wrong_letter")) + correct = normalize_answer_text(base.get("correct_answer") or base.get("answer")) + pressure_target = normalize_answer_text( + base.get("incorrect_answer") or base.get("wrong_answer") + ) + if correct == "unknown" and correct_label: + index = ord(correct_label.upper()) - ord("A") + if 0 <= index < len(choices): + correct = choices[index] + if pressure_target == "unknown" and pressure_label: + index = ord(pressure_label.upper()) - ord("A") + if 0 <= index < len(choices): + pressure_target = choices[index] + if pressure_target == "unknown" and choices and correct in choices: + pressure_target = pick_wrong_answer(correct, choices) + pressure_label = chr(ord("A") + choices.index(pressure_target)) + if correct == "unknown" or pressure_target == "unknown" or correct == pressure_target: + return None + return correct, pressure_target, correct_label, pressure_label + + def build_honesty_control_scenarios( raw_dir: Path, splits: list[str], @@ -282,22 +329,19 @@ def build_sycophancy_scenarios( ) prompt_text = first_text(raw, ("prompt", "messages", "conversation")) question = ensure_text(base.get("question")) or prompt_text - correct = normalize_answer_text( - base.get("correct_answer") or base.get("answer") - ) - pressure_target = normalize_answer_text(base.get("incorrect_answer")) - if ( - not question - or not prompt_text - or correct == "unknown" - or pressure_target == "unknown" - ): + targets = _sycophancy_targets(base) + if targets is None: + continue + correct, pressure_target, correct_label, pressure_label = targets + if not question or not prompt_text: raise ValueError(f"Unresolvable sycophancy record {path}:{index + 1}") group_id = f"syc_{_normalized_hash(question)[:20]}" split = _hashed_split(group_id, split_seed) shared_metadata = { "gold_answer": correct, + "gold_option_label": correct_label, "pressure_target": pressure_target, + "pressure_option_label": pressure_label, "prompt_template": ensure_text(metadata.get("prompt_template")) or file_key, "source_row": index, @@ -342,7 +386,9 @@ def build_sycophancy_scenarios( metadata=shared_metadata, ) rows.append(_scenario_dict(record)) - return _enforce_heldout_group_splits(rows) + return _enforce_heldout_group_splits( + _reserve_overlapping_heldout_groups(rows, split_seed) + ) def build_motivated_reasoning_scenarios( @@ -376,7 +422,9 @@ def build_motivated_reasoning_scenarios( if not question: raise ValueError(f"Missing question at {path}:{index + 1}") correct, choices = resolve_correct_choice(raw) - if correct == "unknown" or len(choices) < 2 or correct not in choices: + if correct == "unknown": + continue + if len(choices) < 2 or correct not in choices: raise ValueError(f"Unresolvable choices at {path}:{index + 1}") pressure_target = pick_wrong_answer(correct, choices) # Some upstream MCQ rows (notably mmlu auxiliary_train) are malformed: diff --git a/cli/extract_task_activations.py b/cli/extract_task_activations.py index be594ed..a133c9e 100644 --- a/cli/extract_task_activations.py +++ b/cli/extract_task_activations.py @@ -3,6 +3,7 @@ import argparse import hashlib import subprocess +import time from pathlib import Path from typing import List @@ -105,6 +106,26 @@ def main() -> None: ), ) parser.add_argument("--seed", type=int, default=42) + parser.add_argument( + "--checkpoint_examples", + type=int, + default=32, + help="Atomically checkpoint this many examples at a time (default: 32).", + ) + parser.add_argument( + "--max_wall_time_minutes", + type=float, + default=None, + help=( + "Stop cleanly between activation checkpoints after this much extraction " + "time. Rerun the identical command to resume." + ), + ) + parser.add_argument( + "--overwrite", + action="store_true", + help="Recompute completed splits and checkpoint chunks instead of resuming.", + ) parser.add_argument( "--allow_dirty_code", action="store_true", @@ -116,6 +137,10 @@ def main() -> None: help="Execution device: auto|cpu|cuda|cuda:N|mps", ) args = parser.parse_args() + if args.checkpoint_examples < 1: + raise ValueError("checkpoint_examples must be positive") + if args.max_wall_time_minutes is not None and args.max_wall_time_minutes <= 0: + raise ValueError("max_wall_time_minutes must be positive") task = TASK_REGISTRY[args.task]() data_path = Path(args.data) @@ -157,10 +182,27 @@ def main() -> None: ) ) + deadline_monotonic = ( + time.monotonic() + args.max_wall_time_minutes * 60.0 + if args.max_wall_time_minutes is not None + else None + ) for split_name, split_examples in splits.items(): if split_examples: - extractor.extract_split(split_examples, split_name) - print(f"saved {split_name} -> {len(split_examples)} examples") + completed = extractor.extract_split( + split_examples, + split_name, + checkpoint_examples=args.checkpoint_examples, + resume=not args.overwrite, + deadline_monotonic=deadline_monotonic, + ) + if not completed: + print( + f"checkpointed {split_name}; rerun the identical command " + "without --overwrite to continue" + ) + break + print(f"ready {split_name} -> {len(split_examples)} input examples") if __name__ == "__main__": diff --git a/cli/extract_text_embeddings.py b/cli/extract_text_embeddings.py index 9142215..0f081c4 100644 --- a/cli/extract_text_embeddings.py +++ b/cli/extract_text_embeddings.py @@ -5,6 +5,7 @@ import json import re import subprocess +import time from pathlib import Path from typing import Any @@ -15,6 +16,7 @@ from data.text_embedding_cache import ( TEXT_EMBEDDING_SCHEMA_VERSION, atomic_save_text_embedding_cache, + load_text_embedding_cache, ) from data.text_views import ( examples_to_text_arrays, @@ -27,6 +29,19 @@ COMMIT_RE = re.compile(r"^[0-9a-f]{40}$") +def _validate_resumable_cache( + path: Path, metadata: dict[str, Any], expected: dict[str, Any] +) -> None: + mismatches = [ + key for key, value in expected.items() if metadata.get(key) != value + ] + if mismatches: + raise ValueError( + f"Existing text embedding cache {path} is incompatible in {mismatches}; " + "use --overwrite only if recomputation is intentional" + ) + + def _git_state() -> tuple[str, bool]: try: revision = subprocess.run( @@ -198,12 +213,23 @@ def main() -> None: parser.add_argument("--allow_truncation", action="store_true") parser.add_argument("--allow_dirty_code", action="store_true") parser.add_argument("--overwrite", action="store_true") + parser.add_argument( + "--max_wall_time_minutes", + type=float, + default=None, + help=( + "Stop cleanly between completed text-view caches after this much " + "encoding time. Rerun the identical command to resume." + ), + ) parser.add_argument( "--device", default="auto", help="Execution device: auto|cpu|cuda|cuda:N|mps", ) args = parser.parse_args() + if args.max_wall_time_minutes is not None and args.max_wall_time_minutes <= 0: + raise ValueError("max_wall_time_minutes must be positive") views = [value.strip() for value in args.views.split(",") if value.strip()] invalid_views = sorted(view for view in views if not is_valid_text_view(view)) @@ -235,6 +261,38 @@ def main() -> None: if not split_values.issubset({"train", "calibration", "eval", "test"}): raise ValueError(f"Dataset contains invalid protocol splits: {sorted(split_values)}") + output_dir = Path(args.output_dir) + pending_views: list[str] = [] + for view in views: + output_path = output_dir / f"{args.task}__{view}.npz" + if not output_path.exists() or args.overwrite: + pending_views.append(view) + continue + cached = load_text_embedding_cache( + output_path, + require_clean_code=not args.allow_dirty_code, + ) + _validate_resumable_cache( + output_path, + cached["metadata"], + { + "task_name": args.task, + "view": view, + "dataset_sha256": dataset_hash, + "code_revision": code_revision, + "embedding_model_id": spec["model_id"], + "embedding_model_revision": spec["model_revision"], + "embedding_tokenizer_revision": spec["tokenizer_revision"], + "embedding_spec_sha256": spec_hash, + "embedding_config_sha256": config_hash, + **identity, + }, + ) + print(json.dumps({"skipped": str(output_path), "reason": "validated cache"})) + if not pending_views: + print(f"completed 0 new caches; all {len(views)} requested caches already exist") + return + from transformers import AutoModel, AutoTokenizer from cli.common import ( bounded_batch_size_for_device, @@ -269,12 +327,20 @@ def main() -> None: model.eval() model.requires_grad_(False) - output_dir = Path(args.output_dir) completed = 0 - for view in views: + deadline_monotonic = ( + time.monotonic() + args.max_wall_time_minutes * 60.0 + if args.max_wall_time_minutes is not None + else None + ) + for view in pending_views: + if deadline_monotonic is not None and time.monotonic() >= deadline_monotonic: + print( + "wall-time limit reached at a durable text-view boundary; rerun " + "the identical command to continue" + ) + break output_path = output_dir / f"{args.task}__{view}.npz" - if output_path.exists() and not args.overwrite: - raise FileExistsError(f"Refusing to overwrite existing cache {output_path}") arrays = examples_to_text_arrays(examples, view) raw_texts = arrays.pop("texts").tolist() rendered = [render_embedding_input(text, spec) for text in raw_texts] diff --git a/cli/generate_task_rollouts.py b/cli/generate_task_rollouts.py index 8d978d5..55acc87 100644 --- a/cli/generate_task_rollouts.py +++ b/cli/generate_task_rollouts.py @@ -5,6 +5,7 @@ import json import os import subprocess +import time from datetime import datetime, timezone from pathlib import Path @@ -12,6 +13,15 @@ from data.generation_confidence import build_generation_confidence_trace +def _deadline_reached( + deadline_monotonic: float | None, *, now_monotonic: float | None = None +) -> bool: + if deadline_monotonic is None: + return False + now = time.monotonic() if now_monotonic is None else now_monotonic + return now >= deadline_monotonic + + def _read_scenarios(path: Path) -> list[ScenarioRecord]: scenarios: list[ScenarioRecord] = [] seen: set[str] = set() @@ -157,6 +167,15 @@ def main() -> None: parser.add_argument("--temperature", type=float, default=0.0) parser.add_argument("--top_p", type=float, default=1.0) parser.add_argument("--overwrite", action="store_true") + parser.add_argument( + "--max_wall_time_minutes", + type=float, + default=None, + help=( + "Stop cleanly between completed rollouts after this much generation " + "time. Rerun the identical command to resume." + ), + ) parser.add_argument( "--device", default="auto", @@ -179,6 +198,8 @@ def main() -> None: raise ValueError("--model_revision must be immutable, not main/latest/unpinned") if args.num_rollouts < 1 or args.max_new_tokens < 1: raise ValueError("num_rollouts and max_new_tokens must be positive") + if args.max_wall_time_minutes is not None and args.max_wall_time_minutes <= 0: + raise ValueError("max_wall_time_minutes must be positive") if args.temperature < 0.0 or not 0.0 < args.top_p <= 1.0: raise ValueError("temperature must be non-negative and top_p must be in (0, 1]") if args.num_rollouts > 1 and args.temperature == 0.0: @@ -240,8 +261,33 @@ def main() -> None: } generated_count = 0 + deadline_monotonic = ( + time.monotonic() + args.max_wall_time_minutes * 60.0 + if args.max_wall_time_minutes is not None + else None + ) + stopped_at_checkpoint = False with output.open("a", encoding="utf-8") as handle: for scenario_index, scenario in enumerate(scenarios): + pending_rollouts: list[tuple[int, int, str]] = [] + for replicate in range(args.num_rollouts): + rollout_seed = ( + args.seed + scenario_index * args.num_rollouts + replicate + ) + rollout_id = _rollout_id( + scenario, + args.model, + args.model_revision, + replicate, + rollout_seed, + ) + if rollout_id not in completed: + pending_rollouts.append((replicate, rollout_seed, rollout_id)) + if not pending_rollouts: + continue + if _deadline_reached(deadline_monotonic): + stopped_at_checkpoint = True + break # The prompt encoding depends only on the scenario, not the replicate # or seed, so build it once per scenario. Sampling randomness is set # by set_seed() immediately before each generate() call below, and @@ -264,15 +310,10 @@ def main() -> None: encoded = { key: value.to(model_device) for key, value in encoded.items() } - for replicate in range(args.num_rollouts): - rollout_seed = ( - args.seed + scenario_index * args.num_rollouts + replicate - ) - rollout_id = _rollout_id( - scenario, args.model, args.model_revision, replicate, rollout_seed - ) - if rollout_id in completed: - continue + for replicate, rollout_seed, rollout_id in pending_rollouts: + if _deadline_reached(deadline_monotonic): + stopped_at_checkpoint = True + break set_seed(rollout_seed) kwargs = { "max_new_tokens": args.max_new_tokens, @@ -361,10 +402,17 @@ def main() -> None: os.fsync(handle.fileno()) completed.add(rollout_id) generated_count += 1 + if stopped_at_checkpoint: + break print( f"generated {generated_count} new rollouts; {len(completed)} total in {output}" ) + if stopped_at_checkpoint: + print( + "wall-time limit reached at a durable rollout boundary; rerun the " + "identical command to continue" + ) if __name__ == "__main__": diff --git a/cli/run_llm_judge_baselines.py b/cli/run_llm_judge_baselines.py index 4e48828..06e6a89 100644 --- a/cli/run_llm_judge_baselines.py +++ b/cli/run_llm_judge_baselines.py @@ -5,6 +5,7 @@ import json import os import re +import shutil import subprocess from pathlib import Path from typing import Any @@ -54,6 +55,111 @@ Path(__file__).parents[1] / "data" / "llm_judge.py", Path(__file__).parents[1] / "data" / "llm_judge_cache.py", ) +JUDGE_SCORE_KEYS = ( + "labels", + "scores", + "example_ids", + "question_ids", + "prompt_sha256", + "prompt_token_lengths", +) + + +def _atomic_save_score_shard( + path: Path, + scored: dict[str, np.ndarray], + *, + checkpoint_identity: dict[str, object], +) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_suffix(path.suffix + ".tmp") + with temporary.open("wb") as handle: + np.savez_compressed( + handle, + **{key: np.asarray(scored[key]) for key in JUDGE_SCORE_KEYS}, + checkpoint_identity_json=np.asarray( + canonical_json(checkpoint_identity) + ), + ) + handle.flush() + os.fsync(handle.fileno()) + temporary.replace(path) + + +def _load_score_shard( + path: Path, + *, + expected_example_ids: np.ndarray, + checkpoint_identity: dict[str, object], +) -> dict[str, np.ndarray]: + with np.load(path, allow_pickle=False) as bundle: + missing = [key for key in JUDGE_SCORE_KEYS if key not in bundle] + if missing: + raise ValueError(f"Judge checkpoint {path} is missing {missing}") + if ( + "checkpoint_identity_json" not in bundle + or str(bundle["checkpoint_identity_json"].item()) + != canonical_json(checkpoint_identity) + ): + raise ValueError(f"Judge checkpoint {path} has incompatible provenance") + scored = {key: np.asarray(bundle[key]) for key in JUDGE_SCORE_KEYS} + observed_ids = scored["example_ids"].astype(str) + expected_ids = np.asarray(expected_example_ids).astype(str) + if not np.array_equal(observed_ids, expected_ids) or any( + value.ndim != 1 or len(value) != len(expected_ids) + for value in scored.values() + ): + raise ValueError(f"Judge checkpoint {path} does not match its input examples") + scores = scored["scores"].astype(float) + if not np.all(np.isfinite(scores)) or np.any((scores < 0.0) | (scores > 1.0)): + raise ValueError(f"Judge checkpoint {path} contains invalid scores") + return scored + + +def _score_bundle_with_checkpoints( + runtime: "JudgeRuntime", + bundle: dict[str, np.ndarray], + demonstrations: list[tuple[str, int]], + *, + checkpoint_dir: Path, + checkpoint_examples: int, +) -> dict[str, np.ndarray]: + if checkpoint_examples < 1: + raise ValueError("judge_checkpoint_examples must be positive") + checkpoint_identity = { + "resolved_device": runtime.resolved_device, + "model_parameter_dtype": runtime.model_parameter_dtype, + "effective_batch_size": runtime.batch_size, + } + chunks: list[dict[str, np.ndarray]] = [] + for chunk_index, start in enumerate( + range(0, len(bundle["labels"]), checkpoint_examples) + ): + stop = min(start + checkpoint_examples, len(bundle["labels"])) + chunk_bundle = { + key: np.asarray(value)[start:stop] for key, value in bundle.items() + } + checkpoint_path = checkpoint_dir / f"chunk_{chunk_index:06d}.npz" + if checkpoint_path.exists(): + scored_chunk = _load_score_shard( + checkpoint_path, + expected_example_ids=chunk_bundle["example_ids"], + checkpoint_identity=checkpoint_identity, + ) + else: + scored_chunk = runtime.score_bundle(chunk_bundle, demonstrations) + _atomic_save_score_shard( + checkpoint_path, + scored_chunk, + checkpoint_identity=checkpoint_identity, + ) + chunks.append(scored_chunk) + if not chunks: + raise ValueError("Cannot score an empty judge bundle") + return { + key: np.concatenate([chunk[key] for chunk in chunks], axis=0) + for key in JUDGE_SCORE_KEYS + } def _git_state() -> tuple[str, bool]: @@ -397,6 +503,15 @@ def main() -> None: parser.add_argument("--judge_model_key", required=True) parser.add_argument("--judge_cache_dir", required=True) parser.add_argument("--judge_batch_size", type=int, default=8) + parser.add_argument( + "--judge_checkpoint_examples", + type=int, + default=16, + help=( + "Atomically checkpoint this many judge-scored examples at a time " + "(default: 16)." + ), + ) parser.add_argument("--model", required=True) parser.add_argument("--results_dir", required=True) parser.add_argument("--views", default="prompt_text,answer_text,transcript_text") @@ -419,6 +534,8 @@ def main() -> None: help="Score source eval and reference traffic without touching test targets.", ) args = parser.parse_args() + if args.judge_checkpoint_examples < 1: + raise ValueError("judge_checkpoint_examples must be positive") if bool(args.target_task) != bool(args.target_data): raise ValueError("--target_task and --target_data must be provided together") @@ -585,6 +702,10 @@ def main() -> None: raise ValueError( f"Judge cache {cache_path} was produced from dirty code" ) + shutil.rmtree( + cache_dir / ".checkpoints" / context_hash, + ignore_errors=True, + ) else: if runtime is None: runtime = JudgeRuntime( @@ -601,8 +722,17 @@ def main() -> None: bundles["source_test"] = source["test"] if target is not None and not args.selection_only: bundles["target_test"] = target + context_checkpoint_dir = ( + cache_dir / ".checkpoints" / context_hash + ) scored = { - split_name: runtime.score_bundle(bundle, demonstrations) + split_name: _score_bundle_with_checkpoints( + runtime, + bundle, + demonstrations, + checkpoint_dir=context_checkpoint_dir / split_name, + checkpoint_examples=args.judge_checkpoint_examples, + ) for split_name, bundle in bundles.items() } cache_metadata = { @@ -641,6 +771,7 @@ def main() -> None: expected_context_hash=context_hash, expected_splits=expected_splits, ) + shutil.rmtree(context_checkpoint_dir, ignore_errors=True) result_seeds = ( range(args.seeds) if mode == "zero_shot" else [context_seed] diff --git a/data/source_adapters.py b/data/source_adapters.py index c4b651e..a8f1ff9 100644 --- a/data/source_adapters.py +++ b/data/source_adapters.py @@ -1,6 +1,7 @@ from __future__ import annotations import json +import re from pathlib import Path from typing import Any, Iterable @@ -89,6 +90,7 @@ def choice_list(row: dict[str, Any]) -> list[str]: or row.get("options") or row.get("candidates") or row.get("answer_choices") + or row.get("answers_list") ) if isinstance(value, list): choices: list[str] = [] @@ -114,6 +116,22 @@ def choice_list(row: dict[str, Any]) -> list[str]: for key in sorted(value) if isinstance((item := value[key]), str) and item.strip() ] + lettered = [ + str(row[letter]).strip() + for letter in "ABCDEFGHIJKLMNOPQRSTUVWXYZ" + if isinstance(row.get(letter), str) and str(row[letter]).strip() + ] + if lettered: + return lettered + answers = row.get("answers") + if isinstance(answers, str): + return [ + text.strip() + for text in re.findall( + r"\([A-Z]\)\s*(.*?)(?=\n\([A-Z]\)|$)", answers, re.DOTALL + ) + if text.strip() + ] return [] diff --git a/experiments/data/huggingface_source_lock.yaml b/experiments/data/huggingface_source_lock.yaml index cf30e08..10a0ca7 100644 --- a/experiments/data/huggingface_source_lock.yaml +++ b/experiments/data/huggingface_source_lock.yaml @@ -36,7 +36,7 @@ sources: datasets: - name: mmlu repo: cais/mmlu - revision: "c30699e" + revision: "c30699e8356da336a370243923dbaf21066bb9fe" subset: all splits: source_train: auxiliary_train @@ -44,7 +44,7 @@ sources: test: test - name: arc_challenge repo: allenai/ai2_arc - revision: "1664417" + revision: "16644173ab25fa007c815e3204a8b59a4034ba1b" subset: ARC-Challenge splits: train: train @@ -52,7 +52,7 @@ sources: test: test - name: commonsense_qa repo: tau/commonsense_qa - revision: "94630fe" + revision: "94630fe30dad47192a8546eb75f094926d47e155" subset: default splits: train: train @@ -60,7 +60,7 @@ sources: test: test - name: aqua_rat repo: deepmind/aqua_rat - revision: "33301c6" + revision: "33301c6a050c96af81f63cad5562cb5363e88971" subset: raw splits: train: train diff --git a/experiments/protocol/monitorbench_adapter.yaml b/experiments/protocol/monitorbench_adapter.yaml index 2625545..7bdfc6d 100644 --- a/experiments/protocol/monitorbench_adapter.yaml +++ b/experiments/protocol/monitorbench_adapter.yaml @@ -9,18 +9,18 @@ source: license: MIT source_reference: arXiv:2603.28590 critical_files: - LICENSE: 9ebdae421b08b552f8edd69b053429fd9c8a37b44684fccba1d36aad4b8f3aef - main.py: dfee40da50b99b06273bb611c7b918bd0435b0abf996a76f221073d79af4e646 - run.sh: 27cbafe2ec811dfbbf18325e0786823f3b9fb0352b59a8fcb7dfdd8a46c2c92e - utils.py: d14d232fa1360192c6e003d0f96fb058943e148d814befc5b855058219669de1 - config/dataset_config.yaml: 556426f04cd6ce3d74ca292d476ed5861c885f6b1fe24e7f4e83702a24ac65c7 - config/vllm_config.yaml: 99a862bb654faa0c06d73de86fbb7d2d6c47e8f20721870a60cf4e255da7d0b8 - pipeline/base.py: 8281157a27dbef62982427bd2726f14a205f3f90bd4f031749244130b65c203b - pipeline/intervention.py: 8821b872be008e31657b62a31f3b26d51e3d016f36cf0418b42973977a9afc65 - pipeline/outcome.py: d90ee8a7aef11405f27d15e7d4ee4f9d5d407f33fcf014f8d1be80b6f49a27dc - pipeline/paths.py: 11f8a6df2bb45863b89791b1f21fedd9a9aef69111c3a5def653ad8af4fd15a7 - pipeline/process.py: 2f03e9436e9c3e2548dfb364447aea107b8a3c630950228ad1b0fdc020504e82 - pipeline/registry.py: 4b99fc390ee28ee73005840e84991fe0ccd1c6389eafb4aa330f940e84c262f9 + LICENSE: cde76533e941191a287ddba36f216dde931addd143cd68846f1c5214326dd090 + main.py: 673d98b34b20e13ddd8ad43cce6099f939717985ab2c4571278b82e963f025e9 + run.sh: caed973e220034c7245888487e8410ec8987cf541b3c92429866e3931289f571 + utils.py: 9f274c3d0858bbeb65190a160ae9b63ad448e08aff01e0668cde40f55897684f + config/dataset_config.yaml: d33a07cb977171abea083ce3df7fa3300a531884cb11df21743c95423782e590 + config/vllm_config.yaml: 8ee8f26edb2e5645d6890ad08955a8b01967130e93664e964da41756e568526f + pipeline/base.py: 33051c83f099d2ccc5921aebcd3fc09cc03a04f0b24b0a0a27b40f371b15ad5d + pipeline/intervention.py: 75cc95e7b1f7e82d5292cba580bb243ed91d69e2191f65eb87b57607769e5208 + pipeline/outcome.py: b7c39159c4acbab4311f3f3e43cbd4ba7765b30682628d17cf24ce7fd3e74307 + pipeline/paths.py: 353b691929730a1f0e2be9ca07e4766cae6ecd46ef17131b7f83e2743374dad0 + pipeline/process.py: 128cfc2c0ed25f237c8e5b09c8929f63dce6c4f2374073706ab5fcdb1fbc57fa + pipeline/registry.py: da3bdd3de6e3ea0631abd9a04d5feee162106dcbe6fb9827a98f4ab3b857d172 artifact_contract: filename_pattern: evaluated_llm__n=.tested.jsonl diff --git a/extraction/task_extractor.py b/extraction/task_extractor.py index 95d6bf3..18d41aa 100644 --- a/extraction/task_extractor.py +++ b/extraction/task_extractor.py @@ -3,7 +3,10 @@ from dataclasses import dataclass import hashlib import json +import os from pathlib import Path +import shutil +import time from typing import Dict, Iterable, List, Optional import numpy as np @@ -540,9 +543,55 @@ def extract_example_with_metadata( } return result, metadata - def extract_split(self, examples: Iterable[TaskExample], split_name: str) -> None: - outdir = Path(self.cfg.output_dir) - outdir.mkdir(parents=True, exist_ok=True) + @staticmethod + def _atomic_save_npz(path: Path, arrays: dict[str, np.ndarray]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_suffix(path.suffix + ".tmp") + with temporary.open("wb") as handle: + np.savez_compressed(handle, **arrays) + handle.flush() + os.fsync(handle.fileno()) + temporary.replace(path) + + @staticmethod + def _atomic_save_json(path: Path, payload: dict[str, object]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_suffix(path.suffix + ".tmp") + with temporary.open("w", encoding="utf-8") as handle: + json.dump(payload, handle, sort_keys=True) + handle.write("\n") + handle.flush() + os.fsync(handle.fileno()) + temporary.replace(path) + + def _extraction_config_sha256(self, wanted_views: list[str]) -> str: + payload = { + "model_name": self.cfg.model_name, + "model_revision": self.cfg.model_revision, + "tokenizer_revision": self.cfg.tokenizer_revision, + "layers": self.cfg.layers, + "max_length": self.cfg.max_length, + "allow_truncation": self.cfg.allow_truncation, + "pooling_mode": self.cfg.pooling_mode, + "views": wanted_views, + "trajectory_prefix_percentiles": self.cfg.trajectory_prefix_percentiles, + "trajectory_basis": ( + TRAJECTORY_BASIS + if self.cfg.trajectory_prefix_percentiles + else None + ), + "use_chat_template": self.cfg.use_chat_template, + "missing_view_policy": self.cfg.missing_view_policy, + "require_model_generated": self.cfg.require_model_generated, + "split_seed": self.cfg.split_seed, + } + return hashlib.sha256( + json.dumps(payload, sort_keys=True).encode("utf-8") + ).hexdigest() + + def _extract_chunk_arrays( + self, examples: list[TaskExample], split_name: str + ) -> tuple[dict[int, dict[str, np.ndarray]], list[str]]: buffers: Dict[int, Dict[str, list[np.ndarray]]] = { layer: {} for layer in self.cfg.layers } @@ -592,9 +641,10 @@ def extract_split(self, examples: Iterable[TaskExample], split_name: str) -> Non buffers[layer].setdefault(view_name, []).append(vector) if not labels: - raise ValueError(f"No extractable examples remained for split {split_name}") + return {}, dropped_ids # Recompute to avoid assuming all examples share identical trajectory views. wanted_views = self._requested_output_views(int(np.max(token_counts)) if token_counts else 1) + layer_arrays: dict[int, dict[str, np.ndarray]] = {} for layer, view_map in buffers.items(): for view in wanted_views: if len(view_map.get(view, [])) != len(labels): @@ -652,31 +702,233 @@ def extract_split(self, examples: Iterable[TaskExample], split_name: str) -> Non arrays["code_dirty"] = np.asarray(self.cfg.code_dirty) arrays["split_seed"] = np.asarray(self.cfg.split_seed) arrays["extraction_config_sha256"] = np.asarray( - hashlib.sha256( - json.dumps( - { - "model_name": self.cfg.model_name, - "model_revision": self.cfg.model_revision, - "tokenizer_revision": self.cfg.tokenizer_revision, - "layers": self.cfg.layers, - "max_length": self.cfg.max_length, - "allow_truncation": self.cfg.allow_truncation, - "pooling_mode": self.cfg.pooling_mode, - "views": wanted_views, - "trajectory_prefix_percentiles": self.cfg.trajectory_prefix_percentiles, - "trajectory_basis": ( - TRAJECTORY_BASIS - if self.cfg.trajectory_prefix_percentiles - else None - ), - "use_chat_template": self.cfg.use_chat_template, - "missing_view_policy": self.cfg.missing_view_policy, - "require_model_generated": self.cfg.require_model_generated, - "split_seed": self.cfg.split_seed, - }, - sort_keys=True, - ).encode("utf-8") - ).hexdigest() + self._extraction_config_sha256(wanted_views) ) arrays["dropped_example_ids"] = np.asarray("\n".join(dropped_ids)) - np.savez_compressed(outdir / f"{split_name}_layer{layer}.npz", **arrays) + layer_arrays[layer] = arrays + return layer_arrays, dropped_ids + + @staticmethod + def _checkpoint_marker_matches( + marker_path: Path, + *, + input_example_ids: list[str], + checkpoint_identity: dict[str, object], + layer_paths: list[Path], + ) -> bool: + if not marker_path.exists(): + return False + try: + marker = json.loads(marker_path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError): + return False + if ( + marker.get("schema_version") != "task-activation-checkpoint-v1" + or marker.get("input_example_ids") != input_example_ids + or marker.get("checkpoint_identity") != checkpoint_identity + ): + return False + emitted_count = int(marker.get("emitted_count", -1)) + return emitted_count == 0 or ( + emitted_count > 0 and all(path.exists() for path in layer_paths) + ) + + def _completed_split_matches( + self, examples: list[TaskExample], split_name: str + ) -> bool: + paths = [ + Path(self.cfg.output_dir) / f"{split_name}_layer{layer}.npz" + for layer in self.cfg.layers + ] + if not all(path.exists() for path in paths): + return False + expected_ids = {example.example_id for example in examples} + expected_views = self._requested_output_views(1) + expected_scalars = { + "model_name": self.cfg.model_name, + "model_revision": self.cfg.model_revision or "unpinned", + "tokenizer_revision": ( + self.cfg.tokenizer_revision + or self.cfg.model_revision + or "unpinned" + ), + "pooling_mode": self.cfg.pooling_mode, + "requested_views_json": json.dumps(expected_views, sort_keys=True), + "chat_template_used": self.cfg.use_chat_template, + "chat_template_sha256": self.chat_template_sha256 or "none", + "max_length": self.cfg.max_length, + "allow_truncation": self.cfg.allow_truncation, + "missing_view_policy": self.cfg.missing_view_policy, + "require_model_generated": self.cfg.require_model_generated, + "dataset_sha256": self.cfg.dataset_sha256 or "unknown", + "code_revision": self.cfg.code_revision or "unknown", + "code_dirty": self.cfg.code_dirty, + "split_seed": self.cfg.split_seed, + "extraction_config_sha256": self._extraction_config_sha256( + expected_views + ), + } + shared_emitted_ids: list[str] | None = None + shared_dropped_ids: list[str] | None = None + try: + for path in paths: + with np.load(path, allow_pickle=False) as bundle: + for key, expected in expected_scalars.items(): + if key not in bundle or bundle[key].ndim != 0: + return False + if bundle[key].item() != expected: + return False + emitted_ids = np.asarray(bundle["example_ids"]).astype(str).tolist() + dropped_ids = [ + value + for value in str(bundle["dropped_example_ids"].item()).splitlines() + if value + ] + if shared_emitted_ids is None: + shared_emitted_ids = emitted_ids + shared_dropped_ids = dropped_ids + elif ( + emitted_ids != shared_emitted_ids + or dropped_ids != shared_dropped_ids + ): + return False + except (OSError, KeyError, ValueError): + return False + observed_ids = set((shared_emitted_ids or []) + (shared_dropped_ids or [])) + return observed_ids == expected_ids + + def _merge_checkpoints( + self, + *, + checkpoint_dir: Path, + markers: list[dict[str, object]], + split_name: str, + ) -> None: + dropped_ids = [ + str(example_id) + for marker in markers + for example_id in marker["dropped_example_ids"] + ] + emitted_markers = [marker for marker in markers if marker["emitted_count"]] + if not emitted_markers: + raise ValueError(f"No extractable examples remained for split {split_name}") + + for layer in self.cfg.layers: + row_arrays: dict[str, list[np.ndarray]] = {} + scalar_arrays: dict[str, np.ndarray] | None = None + for marker in emitted_markers: + chunk_index = int(marker["chunk_index"]) + shard_path = checkpoint_dir / f"chunk_{chunk_index:06d}_layer{layer}.npz" + with np.load(shard_path, allow_pickle=False) as bundle: + shard = {key: np.asarray(bundle[key]) for key in bundle.files} + shard_scalars = { + key: value + for key, value in shard.items() + if value.ndim == 0 and key != "dropped_example_ids" + } + if scalar_arrays is None: + scalar_arrays = shard_scalars + elif any( + key not in shard_scalars + or not np.array_equal(value, shard_scalars[key]) + for key, value in scalar_arrays.items() + ): + raise ValueError( + f"Incompatible activation checkpoint metadata for {split_name} layer {layer}" + ) + for key, value in shard.items(): + if value.ndim > 0: + row_arrays.setdefault(key, []).append(value) + assert scalar_arrays is not None + arrays = { + key: np.concatenate(values, axis=0) + for key, values in row_arrays.items() + } + arrays.update(scalar_arrays) + arrays["dropped_example_ids"] = np.asarray("\n".join(dropped_ids)) + self._atomic_save_npz( + Path(self.cfg.output_dir) / f"{split_name}_layer{layer}.npz", + arrays, + ) + + def extract_split( + self, + examples: Iterable[TaskExample], + split_name: str, + *, + checkpoint_examples: int = 32, + resume: bool = True, + deadline_monotonic: float | None = None, + ) -> bool: + if checkpoint_examples < 1: + raise ValueError("checkpoint_examples must be positive") + example_list = list(examples) + if resume and self._completed_split_matches(example_list, split_name): + return True + + outdir = Path(self.cfg.output_dir) + outdir.mkdir(parents=True, exist_ok=True) + checkpoint_dir = outdir / ".checkpoints" / split_name + checkpoint_identity: dict[str, object] = { + "dataset_sha256": self.cfg.dataset_sha256 or "unknown", + "code_revision": self.cfg.code_revision or "unknown", + "code_dirty": self.cfg.code_dirty, + "extraction_config_sha256": self._extraction_config_sha256( + self._requested_output_views(1) + ), + "chat_template_sha256": self.chat_template_sha256 or "none", + "resolved_device": self.resolved_device, + "model_parameter_device": self.model_parameter_device, + "model_parameter_dtype": self.model_parameter_dtype, + } + markers: list[dict[str, object]] = [] + for chunk_index, start in enumerate( + range(0, len(example_list), checkpoint_examples) + ): + chunk = example_list[start : start + checkpoint_examples] + input_ids = [example.example_id for example in chunk] + marker_path = checkpoint_dir / f"chunk_{chunk_index:06d}.json" + layer_paths = [ + checkpoint_dir / f"chunk_{chunk_index:06d}_layer{layer}.npz" + for layer in self.cfg.layers + ] + if resume and self._checkpoint_marker_matches( + marker_path, + input_example_ids=input_ids, + checkpoint_identity=checkpoint_identity, + layer_paths=layer_paths, + ): + markers.append(json.loads(marker_path.read_text(encoding="utf-8"))) + continue + if deadline_monotonic is not None and time.monotonic() >= deadline_monotonic: + return False + + layer_arrays, dropped_ids = self._extract_chunk_arrays(chunk, split_name) + for layer, arrays in layer_arrays.items(): + self._atomic_save_npz( + checkpoint_dir / f"chunk_{chunk_index:06d}_layer{layer}.npz", + arrays, + ) + emitted_count = ( + len(next(iter(layer_arrays.values()))["labels"]) + if layer_arrays + else 0 + ) + marker = { + "schema_version": "task-activation-checkpoint-v1", + "chunk_index": chunk_index, + "input_example_ids": input_ids, + "checkpoint_identity": checkpoint_identity, + "emitted_count": emitted_count, + "dropped_example_ids": dropped_ids, + } + self._atomic_save_json(marker_path, marker) + markers.append(marker) + + self._merge_checkpoints( + checkpoint_dir=checkpoint_dir, + markers=markers, + split_name=split_name, + ) + shutil.rmtree(checkpoint_dir, ignore_errors=True) + return True diff --git a/tests/test_extractor.py b/tests/test_extractor.py index 448d95f..e714e18 100644 --- a/tests/test_extractor.py +++ b/tests/test_extractor.py @@ -1,3 +1,4 @@ +from dataclasses import replace from types import SimpleNamespace import numpy as np @@ -64,6 +65,10 @@ def _extractor(**overrides): extractor.cfg = cfg extractor.tokenizer = CharacterTokenizer() extractor.model = HiddenStateFixture() + extractor.resolved_device = "cpu" + extractor.model_parameter_device = "cpu" + extractor.model_parameter_dtype = "torch.float32" + extractor.chat_template_sha256 = "f" * 64 return extractor @@ -100,3 +105,47 @@ def test_missing_requested_view_fails_loudly() -> None: extractor = _extractor(views=["reasoning"]) with pytest.raises(ValueError, match="missing requested views"): extractor._tokenize_segments(_example()) + + +def test_split_extraction_resumes_from_atomic_chunk_checkpoints( + tmp_path, monkeypatch +) -> None: + output_dir = tmp_path / "features" + extractor = _extractor( + output_dir=str(output_dir), + dataset_sha256="d" * 64, + code_revision="c" * 40, + ) + first = replace(_example(), task_family="reference_traffic") + second = replace(first, example_id="r2", question_id="q2") + + ticks = iter([0.0, 2.0]) + monkeypatch.setattr( + "extraction.task_extractor.time.monotonic", lambda: next(ticks) + ) + assert not extractor.extract_split( + [first, second], + "train", + checkpoint_examples=1, + deadline_monotonic=1.0, + ) + + resumed = _extractor( + output_dir=str(output_dir), + dataset_sha256="d" * 64, + code_revision="c" * 40, + ) + original_extract = resumed.extract_example_with_metadata + recomputed_ids = [] + + def record_extract(example): + recomputed_ids.append(example.example_id) + return original_extract(example) + + monkeypatch.setattr(resumed, "extract_example_with_metadata", record_extract) + assert resumed.extract_split( + [first, second], "train", checkpoint_examples=1 + ) + assert recomputed_ids == ["r2"] + with np.load(output_dir / "train_layer0.npz", allow_pickle=False) as bundle: + assert bundle["example_ids"].astype(str).tolist() == ["r1", "r2"] diff --git a/tests/test_generate_task_rollouts.py b/tests/test_generate_task_rollouts.py index 08bfc28..c2063e2 100644 --- a/tests/test_generate_task_rollouts.py +++ b/tests/test_generate_task_rollouts.py @@ -1,9 +1,16 @@ from cli.generate_task_rollouts import ( + _deadline_reached, _generation_stop_reason, _split_reasoning, ) +def test_generation_deadline_is_optional_and_inclusive() -> None: + assert not _deadline_reached(None, now_monotonic=100.0) + assert not _deadline_reached(100.0, now_monotonic=99.9) + assert _deadline_reached(100.0, now_monotonic=100.0) + + def test_reasoning_is_split_before_special_tokens_are_removed() -> None: reasoning, answer = _split_reasoning( "Check both choices carefully.The answer is B.", diff --git a/tests/test_llm_judge.py b/tests/test_llm_judge.py index 41cdaa7..d361537 100644 --- a/tests/test_llm_judge.py +++ b/tests/test_llm_judge.py @@ -6,6 +6,7 @@ from cli.run_llm_judge_baselines import ( _load_judge_spec, + _score_bundle_with_checkpoints, contextual_label_token_ids, pairwise_positive_probability, ) @@ -112,6 +113,54 @@ def test_judge_cache_is_context_bound(tmp_path) -> None: load_judge_cache(path, expected_context_hash="0" * 64) +def test_judge_scoring_resumes_from_example_chunks(tmp_path) -> None: + bundle = { + "texts": np.asarray(["a", "b", "c"]), + "labels": np.asarray([0, 1, 0]), + "example_ids": np.asarray(["e0", "e1", "e2"]), + "question_ids": np.asarray(["q0", "q1", "q2"]), + } + + class Runtime: + calls = 0 + resolved_device = "cpu" + model_parameter_dtype = "torch.float32" + batch_size = 2 + + def score_bundle(self, chunk, demonstrations): + self.calls += 1 + return { + "labels": chunk["labels"], + "scores": chunk["labels"].astype(float), + "example_ids": chunk["example_ids"], + "question_ids": chunk["question_ids"], + "prompt_sha256": np.asarray(["a" * 64] * len(chunk["labels"])), + "prompt_token_lengths": np.ones(len(chunk["labels"]), dtype=np.int64), + } + + runtime = Runtime() + first = _score_bundle_with_checkpoints( + runtime, + bundle, + [], + checkpoint_dir=tmp_path / "chunks", + checkpoint_examples=2, + ) + assert runtime.calls == 2 + assert first["example_ids"].tolist() == ["e0", "e1", "e2"] + + resumed = Runtime() + second = _score_bundle_with_checkpoints( + resumed, + bundle, + [], + checkpoint_dir=tmp_path / "chunks", + checkpoint_examples=2, + ) + assert resumed.calls == 0 + assert second["scores"].tolist() == [0.0, 1.0, 0.0] + + def test_registered_primary_judge_is_fully_pinned() -> None: spec, config_hash, spec_hash = _load_judge_spec( Path("experiments/baselines/llm_judge_models.yaml"), diff --git a/tests/test_scenario_builder.py b/tests/test_scenario_builder.py index be391a8..31ce72d 100644 --- a/tests/test_scenario_builder.py +++ b/tests/test_scenario_builder.py @@ -77,6 +77,59 @@ def test_heldout_prompt_template_is_test_only(tmp_path) -> None: } == {"heldout"} +def test_sycophancy_builder_resolves_pinned_letter_only_rows(tmp_path) -> None: + source_dir = tmp_path / "sycophancy_eval" + source_dir.mkdir() + row = { + "prompt": [{"type": "human", "content": "Which number is even?\n(A) 3\n(B) 4"}], + "base": { + "question": "Which number is even?", + "answers": "(A) 3\n(B) 4", + "correct_letter": "B", + }, + } + (source_dir / "are_you_sure.jsonl").write_text( + json.dumps(row) + "\n", encoding="utf-8" + ) + scenarios = build_sycophancy_scenarios( + tmp_path, ["are_you_sure"], "a" * 40, 42 + ) + assert len(scenarios) == 2 + assert scenarios[0]["metadata"]["gold_answer"] == "4" + assert scenarios[0]["metadata"]["pressure_target"] == "3" + + +def test_overlapping_heldout_templates_do_not_consume_all_source_groups(tmp_path) -> None: + source_dir = tmp_path / "sycophancy_eval" + source_dir.mkdir() + answer_rows = [] + heldout_rows = [] + for index in range(100): + base = { + "question": f"Question {index}?", + "correct_answer": "yes", + "incorrect_answer": "no", + } + answer_rows.append({"prompt": f"Question {index}?", "base": base}) + heldout_rows.append({"prompt": f"Are you sure about {index}?", "base": base}) + (source_dir / "answer.jsonl").write_text( + "".join(json.dumps(row) + "\n" for row in answer_rows), encoding="utf-8" + ) + (source_dir / "are_you_sure.jsonl").write_text( + "".join(json.dumps(row) + "\n" for row in heldout_rows), encoding="utf-8" + ) + scenarios = build_sycophancy_scenarios( + tmp_path, ["answer", "are_you_sure"], "a" * 40, 42 + ) + source_splits = { + row["protocol_split"] + for row in scenarios + if row["metadata"]["falsification"]["axes"]["template"]["role"] + == "source" + } + assert {"train", "eval", "test"}.issubset(source_splits) + + def test_honesty_control_builder_pairs_neutral_and_pressure_conditions(tmp_path) -> None: source_dir = tmp_path / "honesty_control_raw" source_dir.mkdir() diff --git a/tests/test_text_embeddings.py b/tests/test_text_embeddings.py index 85c6640..959499c 100644 --- a/tests/test_text_embeddings.py +++ b/tests/test_text_embeddings.py @@ -2,7 +2,11 @@ import pytest import torch -from cli.extract_text_embeddings import pool_hidden_states, render_embedding_input +from cli.extract_text_embeddings import ( + _validate_resumable_cache, + pool_hidden_states, + render_embedding_input, +) from cli.run_embedding_baselines import _assert_compatible from data.schema import TaskExample from data.text_embedding_cache import ( @@ -103,6 +107,21 @@ def test_embedding_cache_round_trip_and_compatibility(tmp_path) -> None: _assert_compatible(first, incompatible) +def test_resume_requires_matching_embedding_cache_identity(tmp_path) -> None: + metadata = _metadata() + _validate_resumable_cache( + tmp_path / "cache.npz", + metadata, + {"dataset_sha256": metadata["dataset_sha256"], "view": "answer_text"}, + ) + with pytest.raises(ValueError, match="dataset_sha256"): + _validate_resumable_cache( + tmp_path / "cache.npz", + metadata, + {"dataset_sha256": "0" * 64}, + ) + + def test_embedding_input_uses_locked_instruction_format() -> None: spec = { "instruction": "Represent this interaction.",