File size: 7,055 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 | """Local FP32 comparison of the WSD-schedule runs on the four pinned nonblind pools.
Same scorer and semantics as `evaluate_v5.py`. Two runs: V5-50ep (d512/L8, V4.3 mixture)
and V4.2-50ep (d384/L7, V4.2 mixture), both warmup -> 3e-4 stable to 85% -> cosine to 6e-5.
Epoch 50 is the best complete-LM-BPC checkpoint of each run (primary). Epoch 42 is the
last epoch of the stable phase, named before local IME evaluation as a within-run control:
same weights lineage, about 0.13 higher BPC. It is descriptive only, and must be scored
before epoch 50, which pairs against it. 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 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_v4 import JapaneseLMV4
from vimeml.training.infer_v5 import JapaneseLMV5
from vimeml.training.model_v4 import PARAMETERS as V4_PARAMETERS
from vimeml.training.model_v5 import PARAMETERS as V5_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
DOWNLOAD = ROOT / 'outputs/training-wsd-20261024'
RESULTS = ROOT / 'outputs/evaluation-wsd-20261024'
# The V5 run is the repository root; the V4.2 run is stored under v42-50ep/.
RUNS = {
'v5-50ep': {'prefix': '', 'loader': JapaneseLMV5, 'parameters': V5_PARAMETERS,
'steps': {42: 4516228, 50: 5376440}},
'v42-50ep': {'prefix': 'v42-50ep/', 'loader': JapaneseLMV4, 'parameters': V4_PARAMETERS,
'steps': {42: 2272072, 50: 2704835}},
}
BASELINES = {
'v21': ROOT / 'outputs/evaluation-v4-20261015/v21',
'v42-epoch20': ROOT / 'outputs/evaluation-v42-20261020',
'v43-epoch30': ROOT / 'outputs/evaluation-v43-20261020/epoch30',
'v5-epoch30': ROOT / 'outputs/evaluation-v5-20261021/epoch30',
'v5-10ep': ROOT / 'outputs/evaluation-v5-20261021/epoch10',
}
def read(path):
return json.loads(path.read_text(encoding='utf-8'))
def main():
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument('--run', choices=sorted(RUNS), required=True)
parser.add_argument('--epoch', type=int, choices=[42, 50], required=True)
args = parser.parse_args()
run = RUNS[args.run]
relative = f"{run['prefix']}checkpoints/epoch-{args.epoch}.pt"
checkpoint = DOWNLOAD / 'remote' / relative
step = run['steps'][args.epoch]
out = RESULTS / args.run / f'epoch{args.epoch}'
if out.exists() and any(out.iterdir()):
parser.error('Completed results are never overwritten')
if not torch.cuda.is_available():
parser.error('The WSD comparison uses an FP32 CUDA GPU')
baselines = dict(BASELINES)
if args.epoch == 50:
stable = RESULTS / args.run / 'epoch42'
if read(stable / 'summary.json')['status'] != 'complete':
parser.error('Score epoch 42 of this run first')
baselines['wsd-epoch42'] = stable
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 = run['loader'](checkpoint, ROOT / 'artifacts/tokenizers/v4-a/ja-unigram-8192', 'cuda')
assert lm.metadata['checkpoint_step'] == step
assert lm.model.parameter_count() == run['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'], 'run': args.run,
'checkpoint_sha256': digest, 'checkpoint_step': step, 'checkpoint_epoch': args.epoch,
'selection': 'epoch50 = best complete LM validation BPC (primary); epoch42 = last stable-phase epoch, within-run control named before local IME evaluation, descriptive only',
'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('WSD_LOCAL_EVALUATION_COMPLETE', flush=True)
if __name__ == '__main__':
main()
|