File size: 1,954 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 | import pytest
from pydantic import ValidationError
from dlb.schema import ConditionalSampleRecord
def valid_conditional_record() -> dict[str, object]:
"""Return a hand-authored LM1B-sized conditional publication record."""
prefix = [7] * 64
continuation = [8] * 64
return {
"sample_id": 0,
"prompt_id": 0,
"completion_id": 0,
"source_index": 12,
"prefix_token_ids": prefix,
"continuation_token_ids": continuation,
"reference_token_ids": [9] * 64,
"full_token_ids": prefix + continuation,
"prefix_text": "fixed prefix",
"continuation_text": "generated continuation",
"reference_text": "source continuation",
"full_text": "fixed prefix generated continuation",
"seed": 42,
"generation_seconds": 0.25,
"prefix_exact_match": True,
}
def test_conditional_record_binds_prefix_full_and_suffix_slices() -> None:
"""Catch publication records whose retained prefix or suffix is inconsistent."""
record = ConditionalSampleRecord.model_validate(valid_conditional_record())
assert record.full_token_ids[:64] == record.prefix_token_ids
assert record.full_token_ids[64:] == record.continuation_token_ids
assert record.prefix_exact_match is True
@pytest.mark.parametrize(
("mutation", "message"),
[
(lambda value: value.update(prefix_exact_match=False), "prefix_exact_match"),
(lambda value: value["full_token_ids"].__setitem__(0, 999), "prefix"),
(lambda value: value.update(reference_token_ids=[1] * 63), "64"),
],
)
def test_conditional_record_rejects_invalid_publication(mutation, message: str) -> None:
"""Catch loss of the hard-prefix or aligned-reference publication invariants."""
value = valid_conditional_record()
mutation(value)
with pytest.raises(ValidationError, match=message):
ConditionalSampleRecord.model_validate(value)
|