File size: 5,041 Bytes
f91d419 | 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 | """Evaluate fixed reconstruction masks and the actual interactive editor."""
import argparse
import json
from pathlib import Path
import time
import numpy as np
import torch
from diffuthink.v2.data import digest
from diffuthink.v2.editor import rewrite_span
from diffuthink.v2.inference import load as load_model, denoise
from diffuthink.v2.hybrid import causal_evaluate
from diffuthink.v2.model import PAD, MASK
from .data import load
from .train import span_evaluate
def normalize(text):
return ' '.join(text.casefold().strip(' .!?\n\t').split())
def run(args):
torch.set_num_threads(4)
arrays, manifest = load(args.data)
model, tokenizer = load_model(args.model, args.device)
if digest(Path(args.model)/'tokenizer.json') != manifest['tokenizer_sha256']:
raise ValueError('Benchmark requires the identical frozen tokenizer')
precision = 'bf16' if args.device == 'cuda' else 'fp32'
report = {'model': args.model, 'weights_sha256': digest(Path(args.model)/'model.safetensors'),
'device': args.device, 'precision': precision, 'torch_version': str(torch.__version__),
'data_manifest': manifest, 'split': args.split, 'editor_attempts': 6,
'challenge_sha256': digest(args.challenges), 'seed': 42,
'limitations': 'Exact reference metrics miss valid alternatives. Authored challenges are diagnostic, not a general grammar or coherence score.'}
report['span'] = span_evaluate(model, arrays[args.split], 400, precision)
report['causal'] = causal_evaluate(model, arrays[args.split][0], 400, precision=precision)
print(json.dumps({'span': report['span'], 'causal': report['causal']}), flush=True)
samples = []
array, spans = arrays[args.split]
for index in range(min(64, len(array))):
clean = [int(t) for t in array[index] if t != PAD]
a, b = map(int, spans[index, 0])
masked = clean[:a]+[MASK]*(b-a)+clean[b:]
result = denoise(model, tokenizer, masked, steps=12, temperature=0, seed=42, precision=precision, allow_eos=False)
pred = result['tokens'][a:b]
samples.append({'index': index, 'reference': tokenizer.decode(clean[a:b]),
'prediction': tokenizer.decode(pred), 'exact': pred == clean[a:b],
'context_preserved': result['tokens'][:a] == clean[:a] and result['tokens'][b:] == clean[b:]})
report['iterative_reconstruction'] = {'samples': samples, 'examples': len(samples),
'exact_match': sum(s['exact'] for s in samples)/len(samples),
'context_preservation': sum(s['context_preserved'] for s in samples)/len(samples),
'oracle_length': True, 'note': 'Reference span length supplied; this is not the editor candidate metric.'}
challenges = json.loads(Path(args.challenges).read_text())['cases']
outputs = []
for case in challenges:
result = rewrite_span(model, tokenizer, case['text'], case['start'], case['end'], precision=precision)
accepted = {normalize(x) for x in case['accepted']}
candidates = result['candidates']
outputs.append({**case, 'top1': bool(candidates and normalize(candidates[0]['replacement']) in accepted),
'top3': any(normalize(c['replacement']) in accepted for c in candidates[:3]),
'context_preserved': all(c['text'] == result['prefix']+c['replacement']+result['suffix'] for c in candidates),
'empty': not candidates, 'latency_ms': result['latency_ms'],
'candidates': [{k: v for k, v in c.items() if k != 'trace'} for c in candidates]})
print(json.dumps({'case': case['id'], 'top1': outputs[-1]['top1'], 'top': candidates[0]['replacement'] if candidates else None}), flush=True)
def summarize(rows):
return {key: sum(row[key] for row in rows)/len(rows) for key in ('top1','top3','context_preserved','empty')}
report['editor'] = {'examples': len(outputs), **summarize(outputs),
'median_latency_ms': float(np.median([r['latency_ms'] for r in outputs])),
'by_category': {c: summarize([r for r in outputs if r['category'] == c]) for c in sorted({r['category'] for r in outputs})},
'samples': outputs}
output = Path(args.output)
output.parent.mkdir(parents=True, exist_ok=True)
output.write_text(json.dumps(report, indent=2, ensure_ascii=False), encoding='utf-8')
print(json.dumps({k:v for k,v in report['editor'].items() if k != 'samples'}), flush=True)
if __name__ == '__main__':
p = argparse.ArgumentParser()
p.add_argument('--model', required=True)
p.add_argument('--data', default='data/processed/storypatch-v2')
p.add_argument('--challenges', default='reports/storypatch-v2/challenges.json')
p.add_argument('--output', required=True)
p.add_argument('--device', choices=['cpu','cuda'], default='cuda')
p.add_argument('--split', choices=['validation','test'], default='test')
run(p.parse_args())
|