from pathlib import Path import hashlib import json import pytest from dlb.conditional_prompts import ( build_prompts, load_protocol, select_source_indices, verify_prompts, ) from scripts.build_conditional_prompts import main as build_main PRODUCTION_CONFIG = """\ schema_version: 1 protocol: c64_zs_v1 selection_seed: 42 sampling_seed: 42 prompt_count: 1024 prefix_length: 64 evaluation_continuation_length: 64 diversity_prompt_count: 256 completions_per_diversity_prompt: 5 datasets: lm1b: model_length: 128 owt: model_length: 1024 """ def test_protocol_has_exact_c64_production_contract(tmp_path: Path) -> None: """Catch protocol drift that would change the published C64 benchmark.""" path = tmp_path / "conditional.yaml" path.write_text(PRODUCTION_CONFIG, encoding="utf-8") protocol = load_protocol(path) assert protocol.protocol == "c64_zs_v1" assert (protocol.prompt_count, protocol.prefix_length) == (1024, 64) assert protocol.datasets["owt"].model_length == 1024 def test_selection_is_unique_stable_and_seed_bound() -> None: """Catch selection that is non-deterministic, duplicated, or seed-insensitive.""" first = select_source_indices(4096, 1024, 42) assert first == select_source_indices(4096, 1024, 42) assert first != select_source_indices(4096, 1024, 43) assert len(first) == len(set(first)) == 1024 class TinyValidation: def __init__(self, rows: list[dict[str, list[int]]]) -> None: self.rows = rows def __len__(self) -> int: return len(self.rows) def __getitem__(self, index: int) -> dict[str, list[int]]: return self.rows[index] def production_protocol(): return load_protocol(Path(__file__).parents[1] / "configs" / "conditional.yaml") def install_fake_load_from_disk(monkeypatch, *, rows: int) -> list[str]: observed: list[str] = [] def fake_load_from_disk(path: str) -> TinyValidation: observed.append(path) return TinyValidation( [{"input_ids": [index % 30_522] * 128} for index in range(rows)] ) monkeypatch.setattr("dlb.conditional_prompts.load_from_disk", fake_load_from_disk) return observed def make_prompt_root(tmp_path: Path) -> Path: (tmp_path / "artifacts").mkdir() (tmp_path / "artifacts" / "data.yaml").write_text( """\ schema_version: 1 datasets: lm1b: repo_id: billion-word-benchmark/lm1b revision: source-revision splits: {validation: test} sequence_length: 128 tokenizer: bert-base-uncased models: {bert-base-uncased: tokenizer-revision} """, encoding="utf-8", ) manifests = tmp_path / "data" / "manifests" manifests.mkdir(parents=True) (manifests / "lm1b.json").write_text( json.dumps( { "dataset": "lm1b", "dataset_id": "billion-word-benchmark/lm1b", "tokenizer_id": "bert-base-uncased", "tokenizer_revision": "tokenizer-revision", "sequence_length": 128, "vocabulary": {"size": 30_522}, "packed_sequence_counts": {"validation": 2048}, } ), encoding="utf-8", ) return tmp_path @pytest.fixture def built_prompt_tree(monkeypatch, tmp_path: Path) -> Path: root = make_prompt_root(tmp_path) install_fake_load_from_disk(monkeypatch, rows=2048) build_prompts(root, "lm1b", production_protocol()) return root def test_build_writes_sha_bound_prompt_records(built_prompt_tree: Path) -> None: """Catch prompt publication that loses source alignment or its integrity bindings.""" manifest_path = built_prompt_tree / "data/manifests/conditional-lm1b-c64.json" prompts_path = built_prompt_tree / "data/conditional/lm1b-c64/prompts.jsonl" manifest = json.loads(manifest_path.read_text(encoding="utf-8")) records = [json.loads(line) for line in prompts_path.read_text(encoding="utf-8").splitlines()] assert len(records) == 1024 assert [record["prompt_id"] for record in records] == list(range(1024)) assert all(len(record["prefix_token_ids"]) == len(record["reference_token_ids"]) == 64 for record in records) assert manifest["source_manifest_sha256"] == hashlib.sha256( (built_prompt_tree / "data/manifests/lm1b.json").read_bytes() ).hexdigest() assert manifest["source_row_count"] == 2048 assert manifest["prompt_file_sha256"] == hashlib.sha256(prompts_path.read_bytes()).hexdigest() def test_verify_rejects_prompt_file_tampering(built_prompt_tree: Path) -> None: """Catch a verifier that trusts the prompt file without checking its digest.""" prompts = built_prompt_tree / "data/conditional/lm1b-c64/prompts.jsonl" prompts.write_bytes(prompts.read_bytes() + b"\n") with pytest.raises(ValueError, match="prompt file SHA-256"): verify_prompts(built_prompt_tree, "lm1b", production_protocol()) def test_verify_rejects_duplicate_json_keys_even_with_a_recomputed_digest( built_prompt_tree: Path, ) -> None: """Catch ambiguous prompt JSON that could otherwise hide a changed field.""" prompts = built_prompt_tree / "data/conditional/lm1b-c64/prompts.jsonl" lines = prompts.read_text(encoding="utf-8").splitlines() lines[0] = lines[0].replace('"prompt_id":0,', '"prompt_id":0,"prompt_id":0,', 1) prompts.write_text("\n".join(lines) + "\n", encoding="utf-8") manifest_path = built_prompt_tree / "data/manifests/conditional-lm1b-c64.json" manifest = json.loads(manifest_path.read_text(encoding="utf-8")) manifest["prompt_file_sha256"] = hashlib.sha256(prompts.read_bytes()).hexdigest() manifest_path.write_text(json.dumps(manifest), encoding="utf-8") with pytest.raises(ValueError, match="duplicate JSON key"): verify_prompts(built_prompt_tree, "lm1b", production_protocol()) def test_build_rejects_tokens_outside_source_vocabulary(monkeypatch, tmp_path: Path) -> None: """Catch invalid processed tokens before they enter a prompt artifact.""" root = make_prompt_root(tmp_path) def invalid_rows(path: Path) -> TinyValidation: return TinyValidation([{"input_ids": [30_522] * 128} for _ in range(2048)]) monkeypatch.setattr("dlb.conditional_prompts.load_from_disk", invalid_rows) with pytest.raises(ValueError, match="outside vocabulary"): build_prompts(root, "lm1b", production_protocol()) def test_build_rejects_source_manifest_validation_count_drift(monkeypatch, tmp_path: Path) -> None: """Catch building against rows that no longer match their data-manifest count.""" root = make_prompt_root(tmp_path) install_fake_load_from_disk(monkeypatch, rows=2048) manifest_path = root / "data/manifests/lm1b.json" manifest = json.loads(manifest_path.read_text(encoding="utf-8")) manifest["packed_sequence_counts"]["validation"] = 2047 manifest_path.write_text(json.dumps(manifest), encoding="utf-8") with pytest.raises(ValueError, match="validation row count"): build_prompts(root, "lm1b", production_protocol()) def test_prompt_cli_uses_processed_validation_without_repacking(monkeypatch, tmp_path: Path) -> None: """Catch the build CLI reading a different source or attempting preprocessing.""" make_prompt_root(tmp_path) config = tmp_path / "configs" / "conditional.yaml" config.parent.mkdir() config.write_text(PRODUCTION_CONFIG, encoding="utf-8") observed = install_fake_load_from_disk(monkeypatch, rows=2048) assert build_main(["--root", str(tmp_path), "--dataset", "lm1b"]) == 0 assert observed == [str(tmp_path / "data/processed/lm1b-bert-128/validation")] def test_prompt_builder_passes_string_paths_to_datasets(monkeypatch, tmp_path: Path) -> None: """Catch datasets/fsspec combinations that reject pathlib.Path values.""" root = make_prompt_root(tmp_path) observed: list[str] = [] def fake_load_from_disk(path: str) -> TinyValidation: if not isinstance(path, str): raise TypeError("argument of type 'PosixPath' is not iterable") observed.append(path) return TinyValidation( [{"input_ids": [index % 30_522] * 128} for index in range(2048)] ) monkeypatch.setattr("dlb.conditional_prompts.load_from_disk", fake_load_from_disk) build_prompts(root, "lm1b", production_protocol()) verify_prompts(root, "lm1b", production_protocol()) assert observed == [ str(tmp_path / "data/processed/lm1b-bert-128/validation"), str(tmp_path / "data/processed/lm1b-bert-128/validation"), ]