File size: 9,681 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 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 | """Training-free rescoring of existing FP32 scores: PMI, token-length and engine-score fusion.
Uses only the per-candidate values already stored in completed `scores.jsonl` files
(contextual and context-free suffix logP sums, scored token counts, AzooKey scores);
no model is run and no reference or blind set is read.
Protocol (written to plan.json by --plan-only before any result is computed):
- every method's parameter is selected on Standard development exact Top-1 alone
(ties: smallest parameter magnitude, i.e. closest to the baseline);
- the primary candidate is the method with the best selected Standard score
(ties: fewer parameters); all others are descriptive;
- Legacy, AJIMEE and Historical are held out; Standard is also split by source
(select on one, test on the other);
- every result reports exact Top-1, notation-lenient Top-1 and semantic errors,
paired against the baseline with exact McNemar tests.
"""
import argparse
import datetime
import hashlib
import json
import sys
from math import comb
from pathlib import Path
ROOT = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(ROOT / 'src'))
from vimeml.training.data import write_json
from vimeml.benchmarks.notation import error_class
TUNE = 'standard-development'
HELD_OUT = ['legacy-regression', 'ajimee', 'historical-development']
MODELS = {
'v42-e20': ROOT / 'outputs/evaluation-v42-20261020',
'v5-e30': ROOT / 'outputs/evaluation-v5-20261021/epoch30',
'v21': ROOT / 'outputs/evaluation-v4-20261015/v21',
}
PMI = [0.05, 0.1, 0.15, 0.2, 0.25, 0.3, 0.4, 0.5, 0.6, 0.8, 1.0]
LENGTH = [-1.0, -0.5, -0.25, 0.25, 0.5, 1.0, 1.5, 2.0]
ENGINE = [0.02, 0.05, 0.1, 0.2, 0.3, 0.5, 0.75, 1.0, 1.5, 2.0]
MARGIN = [0.5, 1.0, 2.0, 3.0, 5.0, 8.0]
PLANS = {
1: {'baseline': [()], 'pmi': [(l,) for l in PMI], 'length': [(a,) for a in LENGTH],
'engine': [(b,) for b in ENGINE], 'pmi+engine': [(l, b) for l in PMI for b in ENGINE]},
# Registered after plan 1: engine fusion only where the LM's own top-2 margin is below tau.
2: {'baseline': [()], 'gated-engine': [(b, t) for b in ENGINE for t in MARGIN]},
}
METHODS = PLANS[1]
def plan_for(version):
plan = {
'tune_track': TUNE, 'held_out_tracks': HELD_OUT, 'models': {k: str(v.relative_to(ROOT)) for k, v in MODELS.items()},
'score': {'baseline': 'c', 'pmi': 'c - lambda*f', 'length': 'c + alpha*n', 'engine': 'c + beta*z',
'pmi+engine': 'c - lambda*f + beta*z'},
'symbols': {'c': 'contextual suffix logP sum', 'f': 'context-free logP sum of the same candidate',
'n': 'scored candidate tokens', 'z': 'AzooKey candidate score'},
'grids': {k: v for k, v in PLANS[version].items() if k != 'baseline'},
'selection': 'per method: max Standard exact Top-1, ties -> smallest parameter magnitude; '
'primary = best selected method on Standard, ties -> fewer parameters; others descriptive',
'cross_validation': 'Standard split by source: select on one source, report on the other',
'ties': 'stable sort over the frozen AzooKey candidate order',
'metrics': ['exact Top-1', 'notation-lenient Top-1 (vimeml.benchmarks.notation, conservative)', 'semantic errors'],
'blind_read': False, 'reference_read': False,
}
if version == 2:
plan['plan_version'] = 2
plan['score'] = {'baseline': 'c', 'gated-engine': 'c + beta*z if c(top1) - c(top2) < tau under c, else c'}
plan['follows'] = 'plan 1 primary (ungated engine fusion) failed on held-out AJIMEE; gate designed on Standard only'
return plan
def exact_p(a, b):
n, k = a + b, min(a, b)
return min(1.0, 2 * sum(comb(n, i) for i in range(k + 1)) / 2 ** n) if n else 1.0
def plan_digest(version):
return hashlib.sha256(json.dumps(plan_for(version), sort_keys=True).encode()).hexdigest()
def load(folder, track):
rows = []
with open(folder / track / 'scores.jsonl', encoding='utf-8') as handle:
for line in handle:
row = json.loads(line)
ctx = row['lm_scores']['contextual']['candidates']
free = row['lm_scores']['context_free']['candidates']
assert [c['text'] for c in row['candidates']] == [c['text'] for c in ctx] == [c['text'] for c in free]
row['features'] = [(c['log_probability_sum'], f['log_probability_sum'], c['scored_tokens'], z['score'])
for c, f, z in zip(ctx, free, row['candidates'])]
rows.append(row)
return rows
def score(method, params, c, f, n, z):
if method == 'baseline':
return c
if method == 'pmi':
return c - params[0] * f
if method == 'length':
return c + params[0] * n
if method == 'engine':
return c + params[0] * z
return c - params[0] * f + params[1] * z
def top1(row, method, params):
if method == 'gated-engine':
ctx = sorted((x[0] for x in row['features']), reverse=True)
if len(ctx) < 2 or ctx[0] - ctx[1] >= params[1]:
return top1(row, 'baseline', ())
method, params = 'engine', params[:1]
values = [score(method, params, *x) for x in row['features']]
best = max(range(len(values)), key=lambda i: (values[i], -i))
return row['candidates'][best]['text']
CACHE = {}
def classify(row, choice):
key = (row['id'], choice)
if key not in CACHE:
CACHE[key] = error_class(row, choice)[0]
return CACHE[key]
def evaluate(rows, method, params):
classes = [classify(r, top1(r, method, params)) for r in rows]
return classes, {'exact': classes.count('correct'),
'lenient': classes.count('correct') + classes.count('notation'),
'semantic': classes.count('semantic')}
def paired(base, other):
out = {}
for metric, ok in [('exact', ('correct',)), ('lenient', ('correct', 'notation'))]:
gain = sum(b not in ok and o in ok for b, o in zip(base, other))
loss = sum(b in ok and o not in ok for b, o in zip(base, other))
out[metric] = {'gain': gain, 'loss': loss, 'net': gain - loss, 'exact_mcnemar_p': exact_p(gain, loss)}
return out
def select(rows, method):
best = None
for params in METHODS[method]:
exact = evaluate(rows, method, params)[1]['exact']
key = (exact, -sum(abs(p) for p in params))
if best is None or key > best[0]:
best = (key, params)
return best[1], best[0][0]
def main():
parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
parser.add_argument('--output', type=Path, required=True)
parser.add_argument('--plan-only', action='store_true', help='Write plan.json before computing any result')
parser.add_argument('--plan', type=int, choices=sorted(PLANS), default=1)
args = parser.parse_args()
global METHODS
METHODS = PLANS[args.plan]
out = args.output.resolve()
if args.plan_only:
if out.exists() and any(out.iterdir()):
parser.error('Use an empty output directory for the plan')
out.mkdir(parents=True, exist_ok=True)
write_json(out / 'plan.json', {**plan_for(args.plan), 'plan_sha256': plan_digest(args.plan),
'created_utc': datetime.datetime.now(datetime.timezone.utc).isoformat(),
'script_sha256': hashlib.sha256(Path(__file__).read_bytes()).hexdigest()})
print('plan written', plan_digest(args.plan))
return
plan = json.loads((out / 'plan.json').read_text(encoding='utf-8'))
if plan['plan_sha256'] != plan_digest(args.plan) or (out / 'summary.json').exists():
parser.error('Plan changed after registration, or results already exist')
results = {}
for model, folder in MODELS.items():
data = {track: load(folder, track) for track in [TUNE, *HELD_OUT]}
for track, rows in data.items():
assert all(top1(r, 'baseline', ()) == r['orders']['lm_context_sum'][0] for r in rows), (model, track)
base = {track: evaluate(rows, 'baseline', ()) for track, rows in data.items()}
chosen = {}
for method in METHODS:
params, tune_exact = select(data[TUNE], method)
entry = {'params': list(params), 'tracks': {}}
for track, rows in data.items():
classes, metrics = evaluate(rows, method, params)
entry['tracks'][track] = {**metrics, 'vs_baseline': paired(base[track][0], classes)}
cv = {}
for source in sorted({r['source'] for r in data[TUNE]}):
fit = [r for r in data[TUNE] if r['source'] == source]
test = [r for r in data[TUNE] if r['source'] != source]
cv_params = select(fit, method)[0]
classes, metrics = evaluate(test, method, cv_params)
cv[f'select_{source}'] = {'params': list(cv_params), 'test_n': len(test), **metrics,
'vs_baseline': paired(evaluate(test, 'baseline', ())[0], classes)}
entry['cross_validation'] = cv
chosen[method] = entry
ranked = sorted((m for m in METHODS if m != 'baseline'),
key=lambda m: (-chosen[m]['tracks'][TUNE]['exact'], len(chosen[m]['params'])))
results[model] = {'primary': ranked[0], 'methods': chosen}
print(model, 'primary', ranked[0], chosen[ranked[0]]['params'], flush=True)
write_json(out / 'summary.json', {'plan_sha256': plan['plan_sha256'], 'results': results,
'blind_read': False, 'reference_read': False})
if __name__ == '__main__':
main()
|