#!/usr/bin/env python3 """Batched, greedy, frozen-base/SFT ablations with isolated Scene loading and full denominators.""" from __future__ import annotations import argparse from contextlib import nullcontext import hashlib import json import os from pathlib import Path import sys import time ROOT = Path(__file__).resolve().parents[1] sys.path.insert(0, str(ROOT / 'src')) from vbench_prompts_compile.inference import AdapterRouter, SceneVerifier, parse_output, validate_direction_output from vbench_prompts_compile.training import user_text def sha(path): digest = hashlib.sha256() with Path(path).open('rb') as handle: for block in iter(lambda: handle.read(8 * 1024 * 1024), b''): digest.update(block) return digest.hexdigest() def jobs_for(rows): jobs = [] scene_caps = sorted({r['input']['caption'] for r in rows if r['task'] == 'scene'}) for mode in ('base', 'sft'): for row in rows: caption_modes = ['intact'] if row['task'] == 'scene' and (row['variant'] == 'canonical' or row['source'] == 'snli_test_locations'): caption_modes += ['missing', 'shuffled'] for caption_mode in caption_modes: inp = dict(row['input']) if caption_mode == 'missing': inp['caption'] = '' elif caption_mode == 'shuffled': i = scene_caps.index(inp['caption']) inp['caption'] = scene_caps[(i + 1) % len(scene_caps)] jobs.append({'key': f'{row["id"]}/{mode}/{caption_mode}', 'id': row['id'], 'task': row['task'], 'mode': mode, 'caption_mode': caption_mode, 'input': inp}) return jobs def main(): p = argparse.ArgumentParser(description=__doc__) p.add_argument('--data', default='data/ablation-v1/eval.jsonl') p.add_argument('--config', required=True) p.add_argument('--out', required=True) p.add_argument('--batch-size', type=int, default=8) p.add_argument('--input-budget', type=int, default=1792) args = p.parse_args() import gc import torch torch.set_num_threads(4) torch.manual_seed(20260919) cfg = json.loads(Path(args.config).read_text()) rows = [json.loads(line) for line in Path(args.data).read_text().splitlines()] out = Path(args.out) out.parent.mkdir(parents=True, exist_ok=True) weights = {task: {name: sha(Path(path) / name) for name in ('adapter_model.safetensors', 'adapter_config.json', 'training_config.json')} for task, path in cfg['adapters'].items()} base_weights = {p.name: sha(p) for p in sorted(Path(cfg['model']).glob('model*.safetensors'))} if not base_weights: raise SystemExit('no local safetensors base weights found') plan = {'data_sha256': sha(args.data), 'config': cfg, 'runner_sha256': sha(__file__), 'inference_sha256': sha(ROOT / 'src/vbench_prompts_compile/inference.py'), 'training_sha256': sha(ROOT / 'src/vbench_prompts_compile/training.py'), 'adapter_sha256': weights, 'base_config_sha256': sha(Path(cfg['model']) / 'config.json'), 'base_weight_sha256': base_weights, 'lock_sha256': sha(ROOT / 'uv.lock'), 'tokenizer_sha256': sha(Path(cfg['model']) / 'tokenizer.json'), 'batch_size': args.batch_size, 'input_budget': args.input_budget, 'parse_max_new_tokens': 128, 'scene_max_new_tokens': 8, 'seed': 20260919, 'decoding': 'greedy, non-thinking, bf16, left padding; Scene newline boundary', 'torch_version': torch.__version__, 'gpu': torch.cuda.get_device_name(), 'cuda_visible_devices': os.environ.get('CUDA_VISIBLE_DEVICES'), 'gpu_total_memory_bytes': torch.cuda.get_device_properties(0).total_memory} plan_path = out.with_suffix('.plan.json') if plan_path.exists(): if json.loads(plan_path.read_text()) != plan: raise SystemExit('frozen inference plan mismatch') else: if out.exists(): raise SystemExit('predictions without plan') plan_path.write_text(json.dumps(plan, indent=2) + '\n') complete = {} if out.exists(): for line in out.read_text().splitlines(): r = json.loads(line) if r['key'] in complete: raise SystemExit('duplicate result key') complete[r['key']] = r all_jobs = jobs_for(rows) resource_path = out.with_suffix('.resources.jsonl') with out.open('a') as output, resource_path.open('a') as resources: for system in ('parse', 'scene'): tasks = ['spatial', 'action', 'objects'] if system == 'parse' else ['scene'] pending = [j for j in all_jobs if j['task'] in tasks and j['key'] not in complete] if not pending: continue torch.cuda.reset_peak_memory_stats() loaded_at = time.monotonic() if system == 'parse': engine = AdapterRouter(base_model_path=cfg['model'], adapters={t: cfg['adapters'][t] for t in tasks}) else: engine = SceneVerifier(base_model_path=cfg['model'], adapter_path=cfg['adapters']['scene']) tokenizer, model = engine.tokenizer, engine.model tokenizer.padding_side = 'left' model.eval() config = model.config.to_dict() resources.write(json.dumps({'event': 'load', 'system': system, 'seconds': time.monotonic() - loaded_at, 'parameters_with_loaded_adapters': sum(p.numel() for p in model.parameters()), 'architecture': {k: config.get(k) for k in ('architectures', 'hidden_size', 'intermediate_size', 'num_hidden_layers', 'num_attention_heads', 'num_key_value_heads', 'head_dim', 'vocab_size', 'tie_word_embeddings')}}) + '\n') resources.flush() for task in tasks: for mode in ('base', 'sft'): if system == 'parse': engine.route(task) jobs = [j for j in pending if j['task'] == task and j['mode'] == mode] for j in jobs: user = user_text({'task': task, 'input': j['input']}, spatial_relation_space='four_directions') j['text'] = tokenizer.apply_chat_template([{'role': 'user', 'content': user}], tokenize=False, add_generation_prompt=True, enable_thinking=False) j['input_tokens'] = len(tokenizer.encode(j['text'])) jobs.sort(key=lambda j: (j['input_tokens'], j['key'])) for start in range(0, len(jobs), args.batch_size): batch = jobs[start:start + args.batch_size] overflow = [j for j in batch if j['input_tokens'] > args.input_budget] batch = [j for j in batch if j['input_tokens'] <= args.input_budget] results = [{**{k: v for k, v in j.items() if k not in ('text', 'input')}, 'raw': '', 'value': None, 'errors': ['input_over_budget'], 'seconds_per_item_amortized': None} for j in overflow] if batch: inputs = tokenizer([j['text'] for j in batch], return_tensors='pt', padding=True) inputs.pop('token_type_ids', None) inputs = inputs.to(model.device) torch.cuda.synchronize() t = time.monotonic() new_tokens = 8 if task == 'scene' else 128 kwargs = {'stop_strings': ['\n'], 'tokenizer': tokenizer} if task == 'scene' else {} with torch.inference_mode(), (model.disable_adapter() if mode == 'base' else nullcontext()): generated = model.generate(**inputs, max_new_tokens=new_tokens, do_sample=False, num_beams=1, temperature=None, top_p=None, top_k=None, pad_token_id=tokenizer.pad_token_id, eos_token_id=tokenizer.eos_token_id, **kwargs) torch.cuda.synchronize() elapsed = time.monotonic() - t tokens = generated[:, inputs['input_ids'].shape[1]:] raws = tokenizer.batch_decode(tokens, skip_special_tokens=True) for j, raw, token_ids in zip(batch, raws, tokens.tolist()): value, errors = parse_output(task, raw.strip()) ended = tokenizer.eos_token_id in token_ids or (task == 'scene' and '\n' in raw) if not ended and len(token_ids) >= new_tokens: value, errors = None, tuple(errors) + ('output_budget_exhausted',) if task == 'spatial' and value is not None: value, more = validate_direction_output(value) errors = tuple(errors) + more results.append({**{k: v for k, v in j.items() if k not in ('text', 'input')}, 'raw': raw.strip(), 'value': value, 'errors': list(errors), 'generated_tokens_including_padding': len(token_ids), 'seconds_per_item_amortized': elapsed / len(batch)}) resources.write(json.dumps({'event': 'batch', 'task': task, 'mode': mode, 'items': len(batch), 'seconds': elapsed, 'peak_allocated_bytes': torch.cuda.max_memory_allocated()}) + '\n') resources.flush() for result in results: output.write(json.dumps(result, ensure_ascii=False) + '\n') complete[result['key']] = result output.flush() print(json.dumps({'system': system, 'task': task, 'mode': mode, 'complete': len(complete)}), flush=True) del engine, model, tokenizer gc.collect() torch.cuda.empty_cache() expected = {j['key'] for j in all_jobs} missing, extra = expected - complete.keys(), complete.keys() - expected report = {'planned': len(expected), 'completed': len(complete), 'missing': sorted(missing), 'extra': sorted(extra), 'predictions_sha256': sha(out), 'resources_sha256': sha(resource_path), 'plan_sha256': sha(plan_path), 'complete': not missing and not extra} out.with_suffix('.report.json').write_text(json.dumps(report, indent=2) + '\n') print(json.dumps(report), flush=True) if __name__ == '__main__': main()