File size: 8,613 Bytes
dc72977
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
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"),
    ]