| """Regression tests for all-completion scoring and exact C64 prompt identity.""" |
| import math |
| import io |
| from contextlib import redirect_stdout |
| from unittest.mock import patch |
| from pathlib import Path |
| import sys |
| import unittest |
|
|
| sys.path.insert(0, str(Path(__file__).resolve().parents[1] / 'scripts')) |
| import evaluate_elf_owt_baseline as evaluator |
| from evaluate_elf_owt_baseline import c64_records |
| from run_elf_remaining import Cell |
| from evaluation.conditional_perplexity import compute_conditional_gen_ppl |
|
|
|
|
| class Tokenizer: |
| def decode(self, ids, **kwargs): |
| return ' '.join(map(str, ids)) |
|
|
| def __call__(self, texts, **kwargs): |
| rows = [[int(x) for x in text.split()][:kwargs['max_length']] for text in texts] |
| width = max(map(len, rows)) |
| return {'input_ids': [r + [0] * (width - len(r)) for r in rows], |
| 'attention_mask': [[1] * len(r) + [0] * (width - len(r)) for r in rows]} |
|
|
|
|
| class Scorer: |
| def eval(self): |
| return self |
|
|
| def score_prompt_excluded_batch(self, ids, attention_mask, target_mask): |
| nll, count = 0., 0 |
| for row, mask in zip(ids, target_mask): |
| for token, target in zip(row[1:], mask[1:]): |
| if target: |
| if token not in (2, 3): |
| raise AssertionError('Prompt or post-64 suffix was scored') |
| nll += math.log(2 if token == 2 else 8) |
| count += 1 |
| return nll, count |
|
|
|
|
| class C64Alignment(unittest.TestCase): |
| def test_all_2048_completions_and_only_first_64_response_tokens(self): |
| tok = Tokenizer() |
| prompts = [{'prompt_id': i, 'prefix_token_ids': [9] * 64, |
| 'reference_token_ids': [4] * 64} for i in range(1024)] |
| schedule = [(i, 0) for i in range(1024)] + [(i, c) for c in range(1, 5) for i in range(256)] |
| rows = [{'prompt_id': i, 'completion_id': c, 'input': tok.decode([9] * 64), |
| 'reference': tok.decode([4] * 64)} for i, c in schedule] |
| tokens = [[2 if c == 0 else 3] * 64 + [99] * 896 for i, c in schedule] |
| records = c64_records(rows, prompts, tokens, tok) |
| result = compute_conditional_gen_ppl(records, Scorer(), tok, tok, batch_size=32) |
| self.assertEqual(result.sample_count, 2048) |
| self.assertEqual(result.valid_token_count, 2048 * 64) |
| self.assertAlmostEqual(result.perplexity, 4.0) |
|
|
| def test_plan_selects_only_requested_steps_without_gpu(self): |
| cells = [Cell('owt-prefix', 'c64-text-t5', s, 1024) for s in (1, 2, 4, 8, 16, 32, 1024)] |
| argv = ['evaluate', '--task', 'owt-prefix', '--steps', '4', '8', |
| '--short-response-policy', 'observed', '--plan'] |
| out = io.StringIO() |
| with patch.object(sys, 'argv', argv), patch.object(evaluator, 'manifest', return_value={'assets': {}}), \ |
| patch.object(evaluator, 'sha256', return_value='hash'), \ |
| patch.object(evaluator, 'build_matrix', return_value=(cells, [])), \ |
| patch.object(evaluator, 'reusable_result', return_value=({'run': '/tmp/source'}, [])) as reuse, \ |
| patch.object(evaluator, 'require_server') as server, redirect_stdout(out): |
| evaluator.main() |
| self.assertEqual([c.args[1].steps for c in reuse.call_args_list], [4, 8]) |
| server.assert_not_called() |
| self.assertIn('2 cells', out.getvalue()) |
| self.assertIn('owt-prefix-observed', out.getvalue()) |
|
|
| def test_observed_policy_keeps_short_output_and_weights_actual_tokens(self): |
| tok = Tokenizer() |
| prompts = [{'prompt_id': i, 'prefix_token_ids': [9] * 64, |
| 'reference_token_ids': [4] * 64} for i in range(2)] |
| rows = [{'prompt_id': i, 'completion_id': 0, 'input': tok.decode([9] * 64), |
| 'reference': tok.decode([4] * 64)} for i in range(2)] |
| records = c64_records(rows, prompts, [[2] * 64, [3]], tok, allow_short=True) |
| with self.assertRaisesRegex(ValueError, 'too few'): |
| compute_conditional_gen_ppl(records, Scorer(), tok, tok) |
| result = compute_conditional_gen_ppl(records, Scorer(), tok, tok, |
| allow_short_continuations=True) |
| self.assertEqual(result.sample_count, 2) |
| self.assertEqual(result.valid_token_count, 65) |
| self.assertAlmostEqual(result.perplexity, math.exp((64 * math.log(2) + math.log(8)) / 65)) |
| with self.assertRaisesRegex(ValueError, 'Short response'): |
| c64_records(rows, prompts, [[2] * 64, []], tok, allow_short=True) |
|
|
| def test_short_response_and_mismatched_prompt_are_not_silently_repaired(self): |
| tok = Tokenizer() |
| prompt = {'prompt_id': 0, 'prefix_token_ids': [9] * 64, 'reference_token_ids': [4] * 64} |
| row = {'prompt_id': 0, 'completion_id': 0, 'input': tok.decode([9] * 64), |
| 'reference': tok.decode([4] * 64)} |
| with self.assertRaisesRegex(ValueError, 'Short response'): |
| c64_records([row], [prompt], [[2] * 63], tok) |
| with self.assertRaisesRegex(ValueError, 'input differs'): |
| c64_records([{**row, 'input': 'wrong'}], [prompt], [[2] * 64], tok) |
|
|
|
|
| if __name__ == '__main__': |
| unittest.main() |
|
|
