File size: 6,350 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 | """Local FP32 comparison of V5 epoch checkpoints on the four pinned nonblind pools.
Same scorer and semantics as `evaluate_v43.py`. Epoch 30 is the best complete-LM-BPC
checkpoint (= best-lm.pt); epoch 29 is descriptive only, named by the cluster agent from
dev IME scores, so it is not a selection candidate. Epoch 10 is the controlled-10ep
best-lm, paired with V4.3 controlled-10ep as a size-only comparison. The fixed reference and
blind sets are not read. Each run writes a separate empty output directory.
"""
import argparse
import copy
import datetime
import json
import math
import sys
import time
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.training.infer_v5 import JapaneseLMV5
from vimeml.training.model_v5 import PARAMETERS
from vimeml.posttraining.infer import TRACKS, read_snapshot
from vimeml.benchmarks.evaluate_ajimee import rerank, summarize, metrics
from vimeml.benchmarks.standard_ime import source_metrics
from vimeml.benchmarks.comparison import paired, rows as read_rows
OUT = None
DOWNLOAD = ROOT / 'outputs/training-v5-20261021'
STEPS = {10: 1075287, 29: 3118362, 30: 3225888}
# The controlled 10-epoch run stores its checkpoints under controlled-10ep/.
RUNS = {10: 'controlled-10ep/'}
BASELINES = {
'v21': ROOT / 'outputs/evaluation-v4-20261015/v21',
'v42-epoch20': ROOT / 'outputs/evaluation-v42-20261020',
'v43-epoch30': ROOT / 'outputs/evaluation-v43-20261020/epoch30',
'v43-epoch27': ROOT / 'outputs/evaluation-v43-20261020/epoch27',
'v43-10ep': ROOT / 'outputs/evaluation-v43-20261020/epoch10',
}
def read(path):
return json.loads(path.read_text(encoding='utf-8'))
def main():
global OUT
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument('--output', type=Path, required=True,
help='Separate empty output directory; completed results are never overwritten')
parser.add_argument('--epoch', type=int, choices=sorted(STEPS), required=True)
args = parser.parse_args()
relative = f"{RUNS.get(args.epoch, '')}checkpoints/epoch-{args.epoch}.pt"
checkpoint = DOWNLOAD / 'remote' / relative
step = STEPS[args.epoch]
OUT = args.output.resolve()
if OUT.exists() and any(OUT.iterdir()):
parser.error('Use an empty output directory to preserve prior evaluations')
if not torch.cuda.is_available():
parser.error('The V5 comparison uses an FP32 CUDA GPU')
receipt = read(DOWNLOAD / 'download-receipt.json')
assert receipt['status'] == 'verified'
digest = file_sha(checkpoint)
assert digest == next(r['sha256'] for r in receipt['files'] if r['path'] == relative)
frozen = {track: read_snapshot(track) for track in TRACKS}
torch.set_num_threads(4)
torch.backends.cuda.matmul.allow_tf32 = False
torch.backends.cudnn.allow_tf32 = False
torch.set_float32_matmul_precision('highest')
lm = JapaneseLMV5(checkpoint, ROOT / 'artifacts/tokenizers/v4-a/ja-unigram-8192', 'cuda')
assert lm.metadata['checkpoint_step'] == step
assert lm.model.parameter_count() == PARAMETERS
assert all(torch.isfinite(p).all().item() for p in lm.model.parameters())
OUT.mkdir(parents=True, exist_ok=True)
write_json(OUT / 'plan.json', {
'created_utc': datetime.datetime.now(datetime.timezone.utc).isoformat(),
'repo': receipt['repo_id'], 'revision': receipt['revision'],
'checkpoint_sha256': digest, 'checkpoint_step': step, 'checkpoint_epoch': args.epoch,
'selection': 'epoch30 = best complete LM validation BPC (primary); epoch29 = cluster-agent pick by dev IME, descriptive only; epoch10 = controlled-10ep best-lm, size-only comparison with V4.3 controlled-10ep',
'tracks': {k: {'cases': len(v), 'sha256': TRACKS[k][1]} for k, v in frozen.items()},
'semantics': 'full-vocabulary suffix logP sum; joint context+candidate tokenization; no final EOS, truncation, or tuning; stable ties',
'precision': 'FP32; TF32 disabled', 'blind_read': False, 'test_read': False,
'script_sha256': file_sha(Path(__file__)),
})
write_json(OUT / 'environment.json', {'torch': str(torch.__version__), 'cuda': torch.version.cuda,
'gpu': torch.cuda.get_device_name(), 'parameters': lm.model.parameter_count(), 'model': lm.metadata})
results = {}
with torch.inference_mode():
for track, original in frozen.items():
rows = copy.deepcopy(original)
began = time.perf_counter()
rerank(lm, rows)
dest = OUT / track
dest.mkdir()
(dest / 'scores.jsonl').write_text(''.join(json.dumps(r, ensure_ascii=False, allow_nan=False) + '\n' for r in rows), encoding='utf-8')
result = {'metrics': summarize(rows), 'elapsed_seconds': time.perf_counter() - began,
'source_metrics': source_metrics(rows, metrics) if 'source' in rows[0] else None}
write_json(dest / 'metrics.json', result)
result['comparisons'] = {}
for name, folder in BASELINES.items():
base = read_rows(folder / track / 'scores.jsonl')
assert len(base) == len(rows)
for a, b in zip(base, rows, strict=True):
assert all(a[k] == b[k] for k in ['id', 'query', 'left_context', 'answers', 'candidates'])
assert a['orders']['azookey'] == b['orders']['azookey']
comparison = paired(base, rows)
comparison['paired'].pop('significant_at_0_05', None)
write_json(dest / f'paired-{name}.json', comparison)
result['comparisons'][name] = comparison['paired']
results[track] = result
print(json.dumps({'track': track, 'top1': result['metrics']['all']['lm_context_sum']['top1_correct'],
'comparisons': result['comparisons']}, ensure_ascii=False), flush=True)
for track in TRACKS:
read_snapshot(track)
assert file_sha(checkpoint) == digest
write_json(OUT / 'summary.json', {'status': 'complete', 'results': results, 'blind_read': False, 'test_read': False})
print('V5_LOCAL_EVALUATION_COMPLETE', flush=True)
if __name__ == '__main__':
main()
|