| import math |
|
|
| import pytest |
|
|
| from dlb.schema import ConditionalSampleRecord |
| from evaluation.conditional_perplexity import ( |
| compute_conditional_gen_ppl, |
| conditional_texts, |
| ) |
|
|
|
|
| class SourceTokenizer: |
| def decode(self, token_ids, **kwargs): |
| del kwargs |
| row = list(token_ids) |
| if len(row) == 64 and set(row) == {1}: |
| return "prompt" |
| if len(row) == 64 and set(row) == {2}: |
| return "generated" |
| if len(row) == 64 and set(row) == {3}: |
| return "reference" |
| if len(row) == 128 and row[:64] == [1] * 64 and row[64:] == [2] * 64: |
| return "prompt generated" |
| if len(row) == 960 and set(row) == {2}: |
| return "full generated suffix" |
| raise ValueError(f"unexpected decoded row length {len(row)}") |
|
|
|
|
| class WhitespaceSuffixTokenizer: |
| def decode(self, token_ids, **kwargs): |
| del kwargs |
| row = list(token_ids) |
| if len(row) == 64 and set(row) == {1}: |
| return "prompt" |
| if len(row) == 64 and set(row) == {4}: |
| return " \n\t" |
| if len(row) == 64 and set(row) == {3}: |
| return "reference" |
| if len(row) == 128 and row[:64] == [1] * 64 and row[64:] == [4] * 64: |
| return "prompt \n\t" |
| raise ValueError(f"unexpected decoded row length {len(row)}") |
|
|
|
|
| class ScorerTokenizer: |
| def __call__(self, texts, **kwargs): |
| rows = { |
| "prompt": [9, 8], |
| "prompt generated": [9, 8, 2, 3], |
| } |
| encoded = [rows[text][: kwargs["max_length"]] for text in texts] |
| width = max(map(len, encoded)) |
| return { |
| "input_ids": [row + [0] * (width - len(row)) for row in encoded], |
| "attention_mask": [[1] * len(row) + [0] * (width - len(row)) for row in encoded], |
| } |
|
|
|
|
| class PromptExcludedScorer: |
| def __init__(self) -> None: |
| self.eval_called = False |
|
|
| def eval(self): |
| self.eval_called = True |
| return self |
|
|
| def score_prompt_excluded_batch(self, input_ids, attention_mask, target_mask): |
| del attention_mask |
| total = 0.0 |
| count = 0 |
| for row, targets in zip(input_ids, target_mask, strict=True): |
| for token, target in zip(row[1:], targets[1:], strict=True): |
| if target: |
| assert token in {2, 3} |
| total += math.log(4) |
| count += 1 |
| return total, count |
|
|
|
|
| def record(*, suffix_length: int = 64) -> ConditionalSampleRecord: |
| prefix = [1] * 64 |
| continuation = [2] * suffix_length |
| return ConditionalSampleRecord( |
| sample_id=0, |
| prompt_id=0, |
| completion_id=0, |
| source_index=10, |
| prefix_token_ids=prefix, |
| continuation_token_ids=continuation, |
| reference_token_ids=[3] * 64, |
| full_token_ids=prefix + continuation, |
| prefix_text="prompt", |
| continuation_text="generated", |
| reference_text="reference", |
| full_text="prompt generated", |
| seed=42, |
| generation_seconds=0.0, |
| prefix_exact_match=True, |
| ) |
|
|
|
|
| def test_conditional_ppl_excludes_prompt_tokens_from_nll() -> None: |
| """Catch prompt tokens being included in conditional generative perplexity.""" |
|
|
| result = compute_conditional_gen_ppl( |
| [record()], |
| PromptExcludedScorer(), |
| ScorerTokenizer(), |
| SourceTokenizer(), |
| batch_size=1, |
| ) |
|
|
| assert result.valid_token_count == 2 |
| assert result.perplexity == pytest.approx(4.0) |
|
|
|
|
| def test_conditional_texts_decode_only_evaluation_suffix_for_long_canvas() -> None: |
| """Catch OWT's 960-token generated suffix leaking into main quality metrics.""" |
|
|
| decoded = conditional_texts([record(suffix_length=960)], SourceTokenizer()) |
|
|
| assert decoded[0].generated_suffix == "generated" |
| assert decoded[0].reference_suffix == "reference" |
| assert decoded[0].prefix_and_generated_suffix == "prompt generated" |
|
|
|
|
| def test_conditional_texts_preserves_whitespace_only_suffix() -> None: |
| """Catch valid whitespace token suffixes being rejected as empty text.""" |
|
|
| value = record() |
| value.continuation_token_ids[:] = [4] * 64 |
| value.full_token_ids[:] = value.prefix_token_ids + value.continuation_token_ids |
|
|
| decoded = conditional_texts([value], WhitespaceSuffixTokenizer()) |
|
|
| assert decoded[0].generated_suffix == " \n\t" |
| assert decoded[0].prefix_and_generated_suffix == "prompt \n\t" |
|
|
