| """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()) |
|
|
