File size: 7,432 Bytes
81cffe7 | 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 | """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()
|