| """Development-only candidate-pool checks for the candidate export v2. |
| |
| Compares the imported n-best 20/50 and empty-context pools of Standard v1 |
| development with the original export, then rescores the n-best 50 pool with |
| V2.1 and V4.2 (FP32, same semantics as the frozen comparisons) and reports |
| exact / notation-lenient Top-1 against the recorded n-best 20 scores. |
| Blind and homophone tracks are refused. |
| """ |
| import argparse |
| import copy |
| import json |
| import sys |
| from collections import Counter |
| from pathlib import Path |
|
|
| ROOT = Path(__file__).resolve().parents[2] |
| sys.path.insert(0, str(ROOT / 'src')) |
| import torch |
| from vimeml.training.data import file_sha, write_json |
| from vimeml.benchmarks.standard_ime import load_standard_export |
| from vimeml.benchmarks.evaluate_ajimee import rerank |
| from vimeml.benchmarks.comparison import rows as read_rows |
| from vimeml.benchmarks.notation import error_class |
| from vimeml.benchmarks import statistics |
|
|
| IMPORTED = ROOT / 'artifacts/benchmarks/ime-candidates-v2-imported' |
| ORIGINAL = ROOT / 'artifacts/benchmarks/ime-standard-ja-v1-candidates/development' |
| ALLOWED = ('standard-v1-development-n20', 'standard-v1-development-n50', 'standard-v1-development-nocontext-n20') |
| MODELS = { |
| 'v21': (ROOT / 'artifacts/models/tiny-ja-v2.1-e16k-d320-l6-extend5/best.pt', |
| ROOT / 'artifacts/tokenizers/ja-unigram-16k-v2', |
| ROOT / 'outputs/evaluation-v4-20261015/v21/standard-development/scores.jsonl'), |
| 'v42-e20': (ROOT / 'outputs/training-v42-20261020/remote/checkpoints/best-lm.pt', |
| ROOT / 'artifacts/tokenizers/v4-a/ja-unigram-8192', |
| ROOT / 'outputs/evaluation-v42-20261020/standard-development/scores.jsonl'), |
| } |
| OUTCOMES = {'exact': ('correct',), 'notation_lenient': ('correct', 'notation')} |
|
|
|
|
| def load(name): |
| if name not in ALLOWED: |
| raise ValueError('Only exposed development pools are allowed here') |
| return load_standard_export(IMPORTED / name)[2] |
|
|
|
|
| def texts(row): |
| return [c['text'] for c in row['candidates']] |
|
|
|
|
| def covered(row): |
| return bool(set(texts(row)) & set(row['answers'])) |
|
|
|
|
| def pool_checks(n20, n50, empty): |
| original = load_standard_export(ORIGINAL)[2] |
| by_id = {r['id']: r for r in original} |
| assert [r['id'] for r in n20] == [r['id'] for r in n50] == [r['id'] for r in empty] |
| context = [i for i, r in enumerate(n20) if r['left_context']] |
| gained = [i for i, (a, b) in enumerate(zip(n20, n50)) if not covered(a) and covered(b)] |
| return { |
| 'determinism_n20': { |
| 'cases': len(n20), |
| 'identical_inputs': sum((r['query'], r['left_context']) == (by_id[r['id']]['query'], by_id[r['id']]['left_context']) for r in n20), |
| 'identical_candidate_lists': sum(texts(r) == texts(by_id[r['id']]) for r in n20)}, |
| 'n50': { |
| 'covered_n20': sum(map(covered, n20)), 'covered_n50': sum(map(covered, n50)), |
| 'gained': len(gained), 'lost': sum(covered(a) and not covered(b) for a, b in zip(n20, n50)), |
| 'n20_subset_of_n50': sum(set(texts(a)) <= set(texts(b)) for a, b in zip(n20, n50)), |
| 'same_azookey_top1': sum(texts(a)[0] == texts(b)[0] for a, b in zip(n20, n50)), |
| 'gained_by_source': dict(Counter(n20[i]['source'] for i in gained)), |
| 'mean_pool_size': sum(len(texts(r)) for r in n50) / len(n50)}, |
| 'empty_context_n20': { |
| 'with_context': len(context), |
| 'identical_candidate_lists': sum(texts(n20[i]) == texts(empty[i]) for i in context)}, |
| } |
|
|
|
|
| def rescore(name, n50): |
| checkpoint, tokenizer, _ = MODELS[name] |
| if name == 'v21': |
| from vimeml.training.infer import JapaneseLM as Model |
| else: |
| from vimeml.training.infer_v4 import JapaneseLMV4 as Model |
| lm = Model(checkpoint, tokenizer, 'cuda') |
| rows = copy.deepcopy(n50) |
| with torch.inference_mode(): |
| rerank(lm, rows) |
| del lm |
| torch.cuda.empty_cache() |
| return rows |
|
|
|
|
| def consistency(recorded, scored): |
| """Contextual scores on candidates shared with the recorded n20 run. |
| |
| Sums start after the prefix all candidates share, so a new n50 candidate that |
| shortens it shifts every score in the case by one constant; only the spread of |
| the per-candidate differences within a case can change a ranking. |
| """ |
| spread, prefix_changed, order_kept = 0.0, 0, 0 |
| for a, b in zip(recorded, scored, strict=True): |
| assert a['id'] == b['id'] |
| ca, cb = a['lm_scores']['contextual'], b['lm_scores']['contextual'] |
| old = {c['text']: c['log_probability_sum'] for c in ca['candidates']} |
| new = {c['text']: c['log_probability_sum'] for c in cb['candidates']} |
| diffs = [new[t] - old[t] for t in old if t in new] |
| spread = max(spread, max(diffs) - min(diffs)) |
| prefix_changed += ca['common_prefix_tokens_including_bos'] != cb['common_prefix_tokens_including_bos'] |
| order_kept += [t for t in a['orders']['lm_context_sum'] if t in new] == \ |
| [t for t in b['orders']['lm_context_sum'] if t in old] |
| return {'max_within_case_spread': spread, 'common_prefix_changed': prefix_changed, |
| 'shared_order_identical': order_kept, 'cases': len(recorded)} |
|
|
|
|
| def compare(recorded, scored): |
| labels = statistics.clusters(recorded, template=None) |
| before = [error_class(r, r['orders']['lm_context_sum'][0])[0] for r in recorded] |
| after = [error_class(r, r['orders']['lm_context_sum'][0])[0] for r in scored] |
| result = {'n20': dict(Counter(before)), 'n50': dict(Counter(after))} |
| for metric, ok in OUTCOMES.items(): |
| result[metric] = statistics.paired([c in ok for c in before], [c in ok for c in after], labels) |
| return result |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) |
| parser.add_argument('--output', type=Path, required=True, help='Separate empty output directory') |
| args = parser.parse_args() |
| out = args.output.resolve() |
| if out.exists() and any(out.iterdir()): |
| parser.error('Use an empty output directory') |
| if not torch.cuda.is_available(): |
| parser.error('FP32 CUDA scoring is required') |
| torch.backends.cuda.matmul.allow_tf32 = False |
| torch.backends.cudnn.allow_tf32 = False |
| torch.set_float32_matmul_precision('highest') |
| n20, n50, empty = (load(n) for n in ALLOWED) |
| summary = {'blind_read': False, 'homophone_read': False, 'script_sha256': file_sha(Path(__file__)), |
| 'pools': pool_checks(n20, n50, empty), 'models': {}} |
| print(json.dumps(summary['pools'], ensure_ascii=False), flush=True) |
| out.mkdir(parents=True, exist_ok=True) |
| for name, (checkpoint, _, recorded_path) in MODELS.items(): |
| recorded = read_rows(recorded_path) |
| assert [r['id'] for r in recorded] == [r['id'] for r in n50] |
| scored = rescore(name, n50) |
| (out / f'{name}-n50-scores.jsonl').write_text( |
| ''.join(json.dumps(r, ensure_ascii=False, allow_nan=False) + '\n' for r in scored), encoding='utf-8') |
| entry = {'checkpoint_sha256': file_sha(checkpoint), 'shared_scores': consistency(recorded, scored), |
| **compare(recorded, scored)} |
| summary['models'][name] = entry |
| print(name, json.dumps({k: v for k, v in entry.items() if k != 'checkpoint_sha256'}, ensure_ascii=False), flush=True) |
| write_json(out / 'summary.json', summary) |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
