| |
| import argparse, json, re, time, os, platform, shutil, subprocess, sys, threading |
| from collections import OrderedDict, deque |
| from concurrent.futures import ThreadPoolExecutor |
| from pathlib import Path |
|
|
| import av |
| import numpy as np |
| from datasets import load_dataset |
|
|
| def sec(s): |
| p = [int(x) for x in str(s).split(':')] |
| return p[-1] + 60 * p[-2] + (3600 * p[-3] if len(p) > 2 else 0) |
|
|
| def norm(x): |
| return re.sub(r'\s+', ' ', re.sub(r'<think>.*?</think>', '', str(x or ''), flags=re.S)).strip().lower().strip(' .') |
|
|
| def correct(pred, gold, fmt): |
| p, g = norm(pred), norm(gold) |
| if fmt == 'binary': return (p.startswith('yes') and g.startswith('yes')) or (p.startswith('no') and g.startswith('no')) |
| if fmt in {'number', 'percentage'}: |
| a = re.search(r'-?\d+(?:\.\d+)?', p); b = re.search(r'-?\d+(?:\.\d+)?', g) |
| return bool(a and b and float(a.group()) == float(b.group())) |
| if fmt == 'time': |
| a = re.search(r'\d{1,2}:\d{2}:\d{2}', p); b = re.search(r'\d{1,2}:\d{2}:\d{2}', g) |
| return bool(a and b and a.group() == b.group()) |
| return p == g or (g and g in p) |
|
|
|
|
| def sample_frames(out, n): |
| out = [out[i] for i in np.linspace(0, len(out) - 1, min(n, len(out))).round().astype(int)] |
| return out + [out[0]] if len(out) == 1 else out |
|
|
|
|
| def frames(src, start, end, n=64): |
| def decode(): |
| c = av.open(str(src)); s = c.streams.video[0] |
| c.seek(max(0, int(start / float(s.time_base))), stream=s, any_frame=True) |
| out = [] |
| for index, f in enumerate(c.decode(s)): |
| t = float(f.pts * s.time_base) if f.pts is not None else None |
| if index == 0 and t is not None and t < max(0, start - 60): |
| c.close(); return [] |
| if t is not None and t > end: break |
| image = f.to_ndarray(format='rgb24') |
| if t is None or t >= start: out.append(image) |
| c.close(); return out |
| out = decode() |
| if not out: |
| |
| out = [np.zeros((224, 224, 3), dtype=np.uint8)] * n |
| out = sample_frames(out, n) if out else [] |
| if not out: raise RuntimeError(f'no frames in {src} for [{start},{end}]') |
| arr = np.stack(out) |
| return arr, {'fps': 5.0, 'duration': len(arr) / 5.0, 'total_num_frames': len(arr), |
| 'frames_indices': list(range(len(arr))), 'video_backend': 'opencv', |
| 'do_sample_frames': False} |
|
|
|
|
| def fixed64_frames(src, output_fps, n=64): |
| c = av.open(str(src)); s = c.streams.video[0] |
| out = [f.to_ndarray(format='rgb24') for f in c.decode(s)] |
| c.close() |
| if len(out) != n: |
| raise RuntimeError(f'expected {n} fixed frames in {src}, decoded {len(out)}') |
| fps = float(output_fps) |
| if fps <= 0: |
| raise RuntimeError(f'invalid output_fps={output_fps} for {src}') |
| arr = np.stack(out) |
| return arr, {'fps': fps, 'duration': len(arr) / fps, 'total_num_frames': len(arr), |
| 'frames_indices': list(range(len(arr))), 'video_backend': 'opencv', |
| 'do_sample_frames': False} |
|
|
|
|
| def fixed_absmeta_frames(src, metadata, n=80): |
| c = av.open(str(src)) |
| stream = c.streams.video[0] |
| out = [f.to_ndarray(format='rgb24') for f in c.decode(stream)] |
| c.close() |
| if len(out) != n: |
| raise RuntimeError(f'expected {n} fixed frames in {src}, decoded {len(out)}') |
| if not isinstance(metadata, dict): |
| raise RuntimeError(f'missing video_metadata for {src}') |
| frame_indices = metadata.get('frames_indices') |
| if metadata.get('video_backend') != 'precomputed_absolute_pts' or not isinstance(frame_indices, list): |
| raise RuntimeError(f'invalid absolute video_metadata for {src}') |
| if len(frame_indices) != n: |
| raise RuntimeError(f'expected {n} frame indices in video_metadata for {src}, got {len(frame_indices)}') |
| return np.stack(out), dict(metadata) |
|
|
|
|
| class VideoCache: |
| def __init__(self, root, max_gib): |
| self.root, self.limit = Path(root), int(max_gib * 1024 ** 3) |
| self.root.mkdir(parents=True, exist_ok=True) |
| self.entries, self.size, self.lock = OrderedDict(), 0, threading.Lock() |
| self.hits = self.misses = self.bytes = 0 |
|
|
| def get(self, src): |
| src = Path(src) |
| key = f'{src.parent.parent.name}_{src.parent.name}_{src.name}' |
| dst = self.root / key |
| with self.lock: |
| if dst.exists(): |
| self.hits += 1; self.entries[dst] = dst.stat().st_size; self.entries.move_to_end(dst) |
| return dst |
| self.misses += 1 |
| size = src.stat().st_size |
| while self.entries and self.size + size > self.limit: |
| old, old_size = self.entries.popitem(last=False) |
| old.unlink(missing_ok=True); self.size -= old_size |
| if size > self.limit: |
| return src |
| tmp = dst.with_suffix(dst.suffix + f'.{os.getpid()}.partial') |
| try: |
| shutil.copyfile(src, tmp); os.replace(tmp, dst) |
| self.entries[dst] = size; self.size += size; self.bytes += size |
| return dst |
| except OSError: |
| tmp.unlink(missing_ok=True) |
| return src |
|
|
|
|
| def read_jsonl(path): |
| if not path.exists(): return [] |
| return [json.loads(line) for line in path.read_text().splitlines() if line.strip()] |
|
|
|
|
| def load_rows(data_root): |
| rows = [] |
| for dataset_name, rel in [('heico', 'heico-focus-vqa/data/segment/test.parquet'), ('lapchole', 'lapchole-focus-vqa/data/segment/test.parquet')]: |
| ds = load_dataset('parquet', data_files=str(Path(data_root) / rel), split='train') |
| rows.extend([{**r, '_dataset': dataset_name, '_input_index': len(rows) + i} for i, r in enumerate(ds)]) |
| return rows |
|
|
|
|
| def load_fixed64_rows(dataset_jsonl): |
| path = Path(dataset_jsonl) |
| rows = [] |
| for i, line in enumerate(path.read_text().splitlines()): |
| if not line.strip(): |
| continue |
| r = json.loads(line) |
| if int(r.get('derived_frame_count', 0)) != 64 or int(r.get('verified_frame_count', 0)) != 64: |
| raise ValueError(f'row {i} id={r.get("id")} is not verified fixed64') |
| rows.append({**r, '_dataset': str(r.get('dataset', 'unknown')).replace('-focus-vqa', ''), |
| '_input_index': len(rows), '_fixed64_root': str(path.parent)}) |
| return rows |
|
|
|
|
| def row_id(row): |
| if row.get('id') is not None: |
| return row['id'] |
| metadata = row.get('segment_metadata') or {} |
| if metadata.get('id') is None: |
| raise ValueError('row is missing both id and segment_metadata.id') |
| return metadata['id'] |
|
|
|
|
| def row_answer_format(row): |
| return row.get('answer_format') or (row.get('segment_metadata') or {}).get('answer_format') |
|
|
|
|
| def row_video_name(row): |
| return row.get('video') or (row.get('segment_metadata') or {}).get('video', '') |
|
|
|
|
| def row_timestamp_start(row): |
| return row.get('timestamp_start') or (row.get('segment_metadata') or {}).get('timestamp_start', '00:00:00') |
|
|
|
|
| def load_fixed80_absmeta_requests(dataset_jsonl, nframes=80): |
| path = Path(dataset_jsonl) |
| rows = [] |
| for i, line in enumerate(path.read_text().splitlines()): |
| if not line.strip(): |
| continue |
| r = json.loads(line) |
| qid = row_id(r) |
| if 'answer' in r: |
| raise ValueError(f'row {i} id={qid} unexpectedly contains an answer') |
| messages = r.get('messages') |
| if not isinstance(messages, list) or [m.get('role') for m in messages] != ['system', 'user']: |
| raise ValueError(f'row {i} id={qid} must contain system,user messages only') |
| user_text = messages[1].get('content') |
| if not isinstance(user_text, str) or user_text.count('<video>') != 1: |
| raise ValueError(f'row {i} id={qid} user message must contain exactly one <video> marker') |
| if 'answer' in r or any(m.get('role') == 'assistant' for m in messages): |
| raise ValueError(f'row {i} id={qid} unexpectedly contains answer content') |
| videos = r.get('videos') |
| if not isinstance(videos, list) or len(videos) != 1: |
| raise ValueError(f'row {i} id={qid} must contain exactly one video') |
| video_path = Path(videos[0]) |
| if not video_path.exists(): |
| raise FileNotFoundError(f'row {i} id={qid} missing video: {video_path}') |
| metadata = (r.get('chat_template_kwargs') or {}).get('video_metadata') |
| if not isinstance(metadata, dict): |
| raise ValueError(f'row {i} id={qid} missing chat_template_kwargs.video_metadata') |
| frame_indices = metadata.get('frames_indices') |
| if metadata.get('video_backend') != 'precomputed_absolute_pts' or not isinstance(frame_indices, list): |
| raise ValueError(f'row {i} id={qid} has invalid absolute video_metadata') |
| if len(frame_indices) != nframes: |
| raise ValueError(f'row {i} id={qid} expected {nframes} frame indices, got {len(frame_indices)}') |
| rows.append({**r, 'id': qid, '_dataset': str((r.get('segment_metadata') or {}).get('dataset', 'unknown')).replace('-focus-vqa', ''), |
| '_input_index': len(rows), '_fixed80_absmeta': True}) |
| return rows |
|
|
|
|
| def detect_dataset_format(dataset_jsonl): |
| path = Path(dataset_jsonl) |
| for line in path.read_text().splitlines(): |
| if line.strip(): |
| row = json.loads(line) |
| if 'chat_template_kwargs' in row and 'segment_metadata' in row: |
| return 'fixed80_absmeta_requests' |
| return 'fixed64_with_refs' |
| raise ValueError(f'empty dataset: {dataset_jsonl}') |
|
|
|
|
| def load_dataset_jsonl(dataset_jsonl, dataset_format, nframes): |
| fmt = detect_dataset_format(dataset_jsonl) if dataset_format == 'auto' else dataset_format |
| if fmt == 'fixed80_absmeta_requests': |
| return load_fixed80_absmeta_requests(dataset_jsonl, nframes), fmt |
| if fmt == 'fixed64_with_refs': |
| return load_fixed64_rows(dataset_jsonl), fmt |
| raise ValueError(f'unknown dataset format: {dataset_format}') |
|
|
|
|
| def pending_rows(rows, existing, order='input'): |
| row_by_id = {str(row_id(r)): r for r in rows} |
| ids = [str(row_id(r)) for r in existing] |
| if len(ids) != len(set(ids)): raise ValueError('resume file contains duplicate IDs') |
| if set(ids) - row_by_id.keys(): raise ValueError('resume file contains IDs outside this split') |
| pending = [r for r in rows if str(row_id(r)) not in set(ids)] |
| if order == 'video_time': |
| pending.sort(key=lambda r: (r['_dataset'], row_video_name(r), sec(row_timestamp_start(r)), r['_input_index'])) |
| return row_by_id, set(ids), pending |
|
|
|
|
| def resolve_checkpoint(checkpoint, train_run_dir=None, require_complete=False, expected_step=None): |
| if checkpoint != 'auto-final': |
| ckpt = Path(checkpoint) |
| else: |
| if not train_run_dir: |
| raise ValueError('--train-run-dir is required when --checkpoint auto-final') |
| roots = [Path(train_run_dir) / 'results' / 'train', Path(train_run_dir)] |
| checkpoints = [] |
| for root in roots: |
| checkpoints.extend(root.glob('**/checkpoint-*')) |
| if not checkpoints: |
| raise FileNotFoundError(f'no checkpoints found under {train_run_dir}') |
| ckpt = max(checkpoints, key=lambda p: int(p.name.split('-')[-1])) |
| required = ['adapter_config.json', 'adapter_model.safetensors', 'trainer_state.json'] |
| missing = [name for name in required if not (ckpt / name).exists()] |
| if missing: |
| raise FileNotFoundError(f'incomplete checkpoint {ckpt}: missing {missing}') |
| state = json.loads((ckpt / 'trainer_state.json').read_text()) |
| step = int(state.get('global_step') or ckpt.name.split('-')[-1]) |
| max_steps = state.get('max_steps') |
| if expected_step is not None and step != expected_step: |
| raise ValueError(f'checkpoint {ckpt} step {step} != expected final step {expected_step}') |
| if require_complete and max_steps is not None and step != int(max_steps): |
| raise ValueError(f'checkpoint {ckpt} is not final: global_step={step}, max_steps={max_steps}') |
| return ckpt, {'global_step': step, 'epoch': state.get('epoch'), 'max_steps': max_steps} |
|
|
|
|
| def make_messages(row): |
| if row.get('_fixed80_absmeta'): |
| messages = [] |
| for message in row['messages']: |
| if message['role'] == 'system': |
| messages.append({'role': 'system', 'content': message['content']}) |
| elif message['role'] == 'user': |
| text = message['content'].replace('<video>', '', 1).strip() |
| messages.append({'role': 'user', 'content': [{'type': 'video'}, {'type': 'text', 'text': text}]}) |
| return messages |
| return [{'role': 'user', 'content': [{'type': 'video'}, {'type': 'text', 'text': row['question']}]}] |
|
|
|
|
| def strip_prediction(text): |
| return re.sub(r'<think>.*?</think>', '', str(text), flags=re.S).strip() |
|
|
|
|
| def validate_rows(rows, nframes, dataset_format, decode_samples=0): |
| ids = [str(row_id(r)) for r in rows] |
| if len(ids) != len(set(ids)): |
| raise ValueError('dataset contains duplicate request IDs') |
| missing_videos = 0 |
| frame_lengths = {} |
| backends = {} |
| user_video_markers = {} |
| message_role_patterns = {} |
| system_prompts = set() |
| for r in rows: |
| if dataset_format == 'fixed80_absmeta_requests': |
| videos = r.get('videos') or [] |
| missing_videos += sum(1 for v in videos if not Path(v).exists()) |
| roles = tuple(m.get('role') for m in (r.get('messages') or [])) |
| message_role_patterns[str(roles)] = message_role_patterns.get(str(roles), 0) + 1 |
| system_prompts.add((r.get('messages') or [{}])[0].get('content')) |
| user_content = (r.get('messages') or [{}, {}])[1].get('content', '') |
| markers = user_content.count('<video>') if isinstance(user_content, str) else -1 |
| user_video_markers[markers] = user_video_markers.get(markers, 0) + 1 |
| metadata = (r.get('chat_template_kwargs') or {}).get('video_metadata') or {} |
| frame_lengths[len(metadata.get('frames_indices') or [])] = frame_lengths.get(len(metadata.get('frames_indices') or []), 0) + 1 |
| backends[metadata.get('video_backend')] = backends.get(metadata.get('video_backend'), 0) + 1 |
| elif dataset_format == 'fixed64_with_refs': |
| src = Path(r['_fixed64_root']) / r['derived_video'] |
| missing_videos += 0 if src.exists() else 1 |
| for r in rows[:max(0, decode_samples)]: |
| if dataset_format == 'fixed80_absmeta_requests': |
| fixed_absmeta_frames(Path(r['videos'][0]), (r.get('chat_template_kwargs') or {})['video_metadata'], nframes) |
| elif dataset_format == 'fixed64_with_refs': |
| fixed64_frames(Path(r['_fixed64_root']) / r['derived_video'], r['output_fps'], nframes) |
| return {'rows': len(rows), 'unique_ids': len(set(ids)), 'duplicate_ids': len(ids) - len(set(ids)), |
| 'missing_video_files': missing_videos, 'frames_indices_len_counts': frame_lengths, |
| 'video_metadata_backend_counts': backends, 'message_role_patterns': message_role_patterns, |
| 'unique_system_prompts': len(system_prompts), 'user_video_marker_counts': user_video_markers, |
| 'decode_checked_rows': max(0, decode_samples)} |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument('--model', required=True); ap.add_argument('--checkpoint') |
| ap.add_argument('--train-run-dir') |
| ap.add_argument('--require-complete-checkpoint', action='store_true') |
| ap.add_argument('--expected-checkpoint-step', type=int) |
| ap.add_argument('--data-root'); ap.add_argument('--dataset-jsonl'); ap.add_argument('--output', required=True) |
| ap.add_argument('--dataset-format', choices=('auto', 'fixed64_with_refs', 'fixed80_absmeta_requests'), default='auto') |
| ap.add_argument('--limit', type=int, default=0); ap.add_argument('--shard-index', type=int, default=0) |
| ap.add_argument('--num-shards', type=int, default=1); ap.add_argument('--nframes', type=int, default=64) |
| ap.add_argument('--tensor-parallel-size', type=int, default=4) |
| ap.add_argument('--batch-size', type=int, default=2) |
| ap.add_argument('--max-model-len', type=int, default=16384) |
| ap.add_argument('--max-new-tokens', type=int, default=64) |
| ap.add_argument('--gpu-memory-utilization', type=float, default=0.9) |
| ap.add_argument('--max-num-seqs', type=int) |
| ap.add_argument('--mm-processor-cache-gb', type=float) |
| ap.add_argument('--lora-name', default='segmentqa_tr005') |
| ap.add_argument('--max-lora-rank', type=int, default=64) |
| ap.add_argument('--disable-lora', action='store_true', |
| help='Run a fully merged model without dynamic vLLM LoRA activation.') |
| ap.add_argument('--merged-source-checkpoint', |
| help='Source adapter checkpoint used to create a fully merged model.') |
| ap.add_argument('--merge-report', |
| help='Merge verification report path for a fully merged model.') |
| ap.add_argument('--disable-tower-connector-lora', action='store_true') |
| ap.add_argument('--enforce-eager', action='store_true', |
| help='Disable torch.compile/CUDAGraphs for vLLM startup paths that are incompatible with multimodal LoRA.') |
| ap.add_argument('--resume', action='store_true') |
| ap.add_argument('--resume-from') |
| ap.add_argument('--execution-order', choices=('input', 'video_time'), default='input') |
| ap.add_argument('--decode-workers', type=int, default=4) |
| ap.add_argument('--prefetch-batches', type=int, default=1) |
| ap.add_argument('--video-cache-dir') |
| ap.add_argument('--video-cache-gib', type=float, default=64) |
| ap.add_argument('--results-subdir', action='store_true') |
| ap.add_argument('--run-id', default='unknown'); ap.add_argument('--issue', action='append', default=[]) |
| ap.add_argument('--dry-run', action='store_true') |
| ap.add_argument('--dry-run-decode-samples', type=int, default=0) |
| args = ap.parse_args() |
| if bool(args.data_root) == bool(args.dataset_jsonl): |
| ap.error('provide exactly one of --data-root or --dataset-jsonl') |
| if not args.disable_lora and not args.checkpoint: |
| ap.error('--checkpoint is required unless --disable-lora is set for a fully merged model') |
| out = Path(args.output); out.mkdir(parents=True, exist_ok=True) |
| results_dir = out / 'results' if args.results_subdir else out |
| results_dir.mkdir(parents=True, exist_ok=True) |
| if args.checkpoint: |
| ckpt, ckpt_info = resolve_checkpoint( |
| args.checkpoint, args.train_run_dir, args.require_complete_checkpoint, args.expected_checkpoint_step) |
| else: |
| ckpt, ckpt_info = None, {} |
| if args.dataset_jsonl: |
| rows, dataset_format = load_dataset_jsonl(args.dataset_jsonl, args.dataset_format, args.nframes) |
| else: |
| rows, dataset_format = load_rows(args.data_root), 'source_parquet' |
| if args.limit: rows = rows[:args.limit] |
| rows = rows[args.shard_index::args.num_shards] |
| raw_path = results_dir / 'raw_predictions.jsonl' |
| if args.resume_from and not raw_path.exists(): shutil.copyfile(args.resume_from, raw_path) |
| existing = read_jsonl(raw_path) if args.resume or args.resume_from else [] |
| row_by_id, done, pending = pending_rows(rows, existing, args.execution_order) |
| validation = validate_rows(rows, args.nframes, dataset_format, args.dry_run_decode_samples) |
| if args.dry_run: |
| report = {'status': 'dry_run_passed', 'run_id': args.run_id, 'dataset_format': dataset_format, |
| 'validation': validation, 'checkpoint': str(ckpt) if ckpt else None, 'checkpoint_info': ckpt_info, |
| 'model': args.model, 'max_new_tokens': args.max_new_tokens, 'temperature': 0, |
| 'lora': {'enabled': not args.disable_lora, 'max_lora_rank': args.max_lora_rank, |
| 'enable_tower_connector_lora': not args.disable_tower_connector_lora}, |
| 'pending': len(pending), 'completed_at_start': len(done)} |
| (results_dir / 'dry_run_report.json').write_text(json.dumps(report, indent=2)) |
| print(json.dumps(report, indent=2)) |
| return |
| from vllm import LLM, SamplingParams |
| commit = subprocess.run(['git', 'rev-parse', 'HEAD'], text=True, capture_output=True).stdout.strip() or 'unknown' |
| dirty = bool(subprocess.run(['git', 'status', '--porcelain'], text=True, capture_output=True).stdout) |
| runmeta = {'status': 'running', 'run_id': args.run_id, 'run_name': out.name, 'task_type': 'test', |
| 'backend': 'vllm', 'model': args.model, 'checkpoint': str(ckpt) if ckpt else None, 'checkpoint_info': ckpt_info, |
| 'count': len(rows), 'completed_at_start': len(done), 'pending': len(pending), |
| 'nframes': args.nframes, 'batch_size': args.batch_size, 'max_model_len': args.max_model_len, |
| 'max_new_tokens': args.max_new_tokens, 'dataset_format': dataset_format, 'validation': validation, |
| 'lora': {'enabled': not args.disable_lora, 'name': args.lora_name, 'max_lora_rank': args.max_lora_rank, |
| 'enable_tower_connector_lora': not args.disable_tower_connector_lora}, |
| 'merged_source_checkpoint': args.merged_source_checkpoint, |
| 'merge_report': args.merge_report, |
| 'enforce_eager': args.enforce_eager, |
| 'tensor_parallel_size': args.tensor_parallel_size, 'precision': 'bfloat16', 'temperature': 0, |
| 'execution_order': args.execution_order, 'decode_workers': args.decode_workers, |
| 'prefetch_batches': args.prefetch_batches, 'video_cache_dir': args.video_cache_dir, |
| 'video_cache_gib': args.video_cache_gib, 'data_root': args.data_root, |
| 'dataset_jsonl': args.dataset_jsonl, 'gpu_memory_utilization': args.gpu_memory_utilization, |
| 'max_num_seqs': args.max_num_seqs, |
| 'dataset_split': 'official SegmentQA test', 'dataset_purpose': 'test', |
| 'python': sys.executable, 'host': platform.node(), 'pid': os.getpid(), |
| 'devices': os.environ.get('CUDA_VISIBLE_DEVICES', 'all-visible'), 'argv': sys.argv, |
| 'code_commit': commit, 'dirty': dirty, 'issues': args.issue, 'started_at': time.time(), |
| 'outputs': {'raw_predictions': str(raw_path), 'answer': str(results_dir / 'answer.json'), |
| 'metrics': str(results_dir / 'metrics.json')}} |
| (out / 'run.json').write_text(json.dumps(runmeta, indent=2)) |
| llm_kwargs = {'model': args.model, 'dtype': 'bfloat16', 'tensor_parallel_size': args.tensor_parallel_size, |
| 'trust_remote_code': True, 'limit_mm_per_prompt': {'video': 1}, |
| 'max_model_len': args.max_model_len, |
| 'gpu_memory_utilization': args.gpu_memory_utilization, |
| 'enforce_eager': args.enforce_eager} |
| if not args.disable_lora: |
| llm_kwargs.update(enable_lora=True, |
| max_lora_rank=args.max_lora_rank, |
| enable_tower_connector_lora=not args.disable_tower_connector_lora) |
| if args.max_num_seqs: |
| llm_kwargs['max_num_seqs'] = args.max_num_seqs |
| if args.mm_processor_cache_gb is not None: |
| llm_kwargs['mm_processor_cache_gb'] = args.mm_processor_cache_gb |
| llm = LLM(**llm_kwargs) |
| tok = llm.get_tokenizer(); sampling = SamplingParams(max_tokens=args.max_new_tokens, temperature=0) |
| lora_request = None |
| if not args.disable_lora: |
| from vllm.lora.request import LoRARequest |
| lora_request = LoRARequest(args.lora_name, 1, str(ckpt)) |
| rawf = raw_path.open('a' if existing else 'w') |
| batch_size = max(1, args.batch_size) |
| cache = VideoCache(args.video_cache_dir, args.video_cache_gib) if args.video_cache_dir else None |
| decode_pool = ThreadPoolExecutor(max_workers=max(1, args.decode_workers)) |
| timings = {'queue_wait_seconds': 0.0, 'generate_seconds': 0.0} |
| def build_requests(batch): |
| def make_request(r): |
| if r.get('_fixed80_absmeta'): |
| video = fixed_absmeta_frames(Path(r['videos'][0]), (r.get('chat_template_kwargs') or {})['video_metadata'], args.nframes) |
| elif 'derived_video' in r: |
| src = Path(r['_fixed64_root']) / r['derived_video'] |
| video = fixed64_frames(src, r['output_fps'], args.nframes) |
| else: |
| src = Path(args.data_root) / ('heico-focus-vqa/videos' if r['_dataset'] == 'heico' else 'lapchole-focus-vqa/videos') / r['video'] |
| if cache: src = cache.get(src) |
| video = frames(src, sec(r['timestamp_start']), sec(r['timestamp_end']), args.nframes) |
| prompt = tok.apply_chat_template(make_messages(r), tokenize=False, add_generation_prompt=True, enable_thinking=False) |
| return {'prompt': prompt, 'multi_modal_data': {'video': video}} |
| return [decode_pool.submit(make_request, r) for r in batch] |
| batches = [pending[i:i + batch_size] for i in range(0, len(pending), batch_size)] |
| queued, next_batch = deque(), 0 |
| def fill_queue(): |
| nonlocal next_batch |
| while next_batch < len(batches) and len(queued) < max(1, args.prefetch_batches): |
| queued.append((batches[next_batch], build_requests(batches[next_batch]))); next_batch += 1 |
| fill_queue() |
| processed = len(done) |
| for batch_index in range(len(batches)): |
| batch, futures = queued.popleft(); wait_start = time.perf_counter() |
| requests = [future.result() for future in futures] |
| timings['queue_wait_seconds'] += time.perf_counter() - wait_start |
| fill_queue() |
| t0 = time.perf_counter() |
| if lora_request is None: |
| results = llm.generate(requests, sampling) |
| else: |
| results = llm.generate(requests, sampling, lora_request=lora_request) |
| elapsed = time.perf_counter() - t0 |
| timings['generate_seconds'] += elapsed |
| for j, (r, result) in enumerate(zip(batch, results)): |
| output = result.outputs[0] |
| pred = output.text; latency = elapsed / len(batch) |
| item = {**r, 'prediction': pred, 'latency': latency, |
| 'finish_reason': getattr(output, 'finish_reason', None), |
| 'generated_token_count': len(getattr(output, 'token_ids', []) or [])} |
| if 'answer' in r: |
| item['correct'] = correct(pred, r['answer'], row_answer_format(r)) |
| rawf.write(json.dumps(item, ensure_ascii=False) + '\n') |
| processed += 1 |
| status = f' correct={item["correct"]}' if 'correct' in item else '' |
| print(f'{processed}/{len(rows)} id={row_id(r)}{status} batch_latency={elapsed:.2f}s', flush=True) |
| rawf.flush() |
| decode_pool.shutdown(wait=True) |
| rawf.close() |
| all_items = read_jsonl(raw_path) |
| all_items.sort(key=lambda item: row_by_id[str(item['id'])]['_input_index']) |
| answers = [{'qID': str(row_id(item)), 'content': strip_prediction(item['prediction']), |
| 'latency': item['latency']} for item in all_items] |
| (results_dir / 'answer.json').write_text(json.dumps(answers, ensure_ascii=False, indent=2)) |
| by_fmt = {} |
| for item in all_items: |
| if 'correct' in item: |
| fmt = row_answer_format(item); by_fmt.setdefault(fmt, [0, 0]); by_fmt[fmt][0] += int(item['correct']); by_fmt[fmt][1] += 1 |
| metrics = {'count': len(answers), 'unique_ids': len({answer['qID'] for answer in answers})} |
| if by_fmt: |
| metrics.update({'correct': sum(v[0] for v in by_fmt.values()), |
| 'accuracy': sum(v[0] for v in by_fmt.values()) / len(answers) if answers else 0, |
| 'by_format': {k: {'correct': v[0], 'total': v[1], 'accuracy': v[0] / v[1]} for k, v in by_fmt.items()}}) |
| (results_dir / 'metrics.json').write_text(json.dumps(metrics, indent=2)) |
| runmeta.update(status='completed', finished_at=time.time(), metrics=metrics, timings=timings, |
| cache=None if not cache else {'hits': cache.hits, 'misses': cache.misses, |
| 'copied_bytes': cache.bytes}) |
| (out / 'run.json').write_text(json.dumps(runmeta, indent=2)) |
|
|
|
|
| if __name__ == '__main__': |
| try: |
| main() |
| except BaseException as exc: |
| if '--output' in sys.argv: |
| status = Path(sys.argv[sys.argv.index('--output') + 1]) / 'run.json' |
| if status.exists(): |
| meta = json.loads(status.read_text()); meta.update(status='failed', finished_at=time.time(), error=repr(exc)) |
| status.write_text(json.dumps(meta, indent=2)) |
| raise |
|
|
