Buckets:
| """Run official Bagel/Anole Zebra-CoT checkpoints on the fixed CV-Bench 20.""" | |
| from __future__ import annotations | |
| import argparse | |
| import json | |
| import time | |
| from pathlib import Path | |
| from PIL import Image | |
| from . import anole, bagel | |
| from .common import evaluation_prompt, fixed_records, read_jsonl, score, summary, write_json | |
| MODELS = {"bagel": bagel, "anole": anole} | |
| def _runner(args: argparse.Namespace): | |
| if args.model == "bagel": | |
| return bagel.BagelRunner(args.model_path), { | |
| "model_repo": bagel.MODEL_REPO, | |
| "model_revision": bagel.MODEL_REVISION, | |
| "inference_repo": bagel.INFERENCE_REPO, | |
| "inference_revision": bagel.INFERENCE_REVISION, | |
| "template": "authors infz_bf16.py system prompt + text/image outer loop", | |
| "generation": bagel.HYPERPARAMETERS | {"max_steps_safety_guard": args.bagel_max_steps}, | |
| } | |
| return anole.AnoleRunner(args.model_path, args.anole_inference_root), { | |
| "model_repo": anole.MODEL_REPO, | |
| "model_revision": anole.MODEL_REVISION, | |
| "inference_repo": anole.INFERENCE_REPO, | |
| "inference_revision": anole.INFERENCE_REVISION, | |
| "template": "authors thinking-with-generated-images inference.py general mode", | |
| "generation": anole.CONFIG, | |
| } | |
| def infer(args: argparse.Namespace) -> None: | |
| import random | |
| import numpy as np | |
| import torch | |
| records = fixed_records(args.data_dir, args.selection) | |
| selected = records[args.shard_id :: args.num_shards] | |
| predictions_dir = args.output_dir / "predictions" | |
| predictions_dir.mkdir(parents=True, exist_ok=True) | |
| shard_path = predictions_dir / f"shard_{args.shard_id:02d}.jsonl" | |
| existing = {row["idx"] for row in read_jsonl(shard_path)} if shard_path.exists() and args.resume else set() | |
| pending = [record for record in selected if record["idx"] not in existing] | |
| runner, provenance = _runner(args) | |
| config = { | |
| "schema": "zebra_official_cvbench20_v1", | |
| "model": args.model, | |
| "data_dir": str(args.data_dir.resolve()), | |
| "model_path": str(args.model_path.resolve()), | |
| "selection": str(args.selection.resolve()), | |
| "shard_id": args.shard_id, | |
| "num_shards": args.num_shards, | |
| "selected": [record["idx"] for record in selected], | |
| "seed_policy": "seed + CV-Bench index", | |
| "seed": args.seed, | |
| "provenance": provenance, | |
| } | |
| write_json(args.output_dir / f"shard_{args.shard_id:02d}_config.json", config) | |
| with shard_path.open("a", encoding="utf-8") as handle: | |
| for position, record in enumerate(pending, start=1): | |
| sample_dir = args.output_dir / "samples" / f"{record['idx']:05d}" | |
| sample_dir.mkdir(parents=True, exist_ok=True) | |
| start = time.monotonic() | |
| prompt = evaluation_prompt(record) | |
| sample_seed = args.seed + record["idx"] | |
| random.seed(sample_seed) | |
| np.random.seed(sample_seed) | |
| torch.manual_seed(sample_seed) | |
| if torch.cuda.is_available(): | |
| torch.cuda.manual_seed_all(sample_seed) | |
| with Image.open(args.data_dir / record["image_path"]) as source: | |
| output = ( | |
| runner.run(prompt, [source.convert("RGB")], sample_dir, max_steps=args.bagel_max_steps) | |
| if args.model == "bagel" | |
| else runner.run(prompt, [source.convert("RGB")], sample_dir) | |
| ) | |
| elapsed = time.monotonic() - start | |
| row = { | |
| "schema": "zebra_official_cvbench20_sample_v1", | |
| "model": args.model, | |
| "idx": record["idx"], | |
| "sample_id": f"cvbench-{record['idx']:05d}", | |
| "type": record["type"], | |
| "task": record["task"], | |
| "source": record["source"], | |
| "question": record["question"], | |
| "choices": record["choices"], | |
| "gold_letter": score(record, "")["gold"], | |
| "prompt": prompt, | |
| "image_path": record["image_path"], | |
| "seconds": elapsed, | |
| "seed": sample_seed, | |
| "provenance": provenance, | |
| **output, | |
| } | |
| row["score"] = score(record, row["assistant_output"]) | |
| (sample_dir / "output.txt").write_text(row["assistant_output"], encoding="utf-8") | |
| write_json(sample_dir / "result.json", row) | |
| handle.write(json.dumps(row, ensure_ascii=False) + "\n") | |
| handle.flush() | |
| print( | |
| "PROGRESS " + json.dumps( | |
| { | |
| "position": position, | |
| "total": len(pending), | |
| "idx": row["idx"], | |
| "correct": row["score"]["correct"], | |
| "images": len(row["generated_images"]), | |
| "stop": row["stop_reason"], | |
| "seconds": round(elapsed, 2), | |
| } | |
| ), | |
| flush=True, | |
| ) | |
| def merge(args: argparse.Namespace) -> None: | |
| records = fixed_records(args.data_dir, args.selection) | |
| expected = [record["idx"] for record in records] | |
| rows = { | |
| row["idx"]: row | |
| for path in sorted((args.output_dir / "predictions").glob("shard_*.jsonl")) | |
| for row in read_jsonl(path) | |
| } | |
| missing = [index for index in expected if index not in rows] | |
| extra = sorted(set(rows) - set(expected)) | |
| if missing or extra: | |
| raise ValueError(f"incomplete fixed CV-Bench output: missing={missing}, extra={extra}") | |
| ordered = [rows[index] for index in expected] | |
| predictions = args.output_dir / "predictions.jsonl" | |
| predictions.write_text("".join(json.dumps(row, ensure_ascii=False) + "\n" for row in ordered), encoding="utf-8") | |
| metrics = summary(ordered) | |
| metrics |= {"model": args.model, "selection_indices": expected} | |
| write_json(args.output_dir / "metrics.json", metrics) | |
| print(json.dumps(metrics, ensure_ascii=False), flush=True) | |
| def parser() -> argparse.ArgumentParser: | |
| result = argparse.ArgumentParser(description=__doc__) | |
| commands = result.add_subparsers(dest="command", required=True) | |
| for name in ("infer", "merge"): | |
| command = commands.add_parser(name) | |
| command.add_argument("--model", choices=tuple(MODELS), required=True) | |
| command.add_argument("--data-dir", type=Path, default=Path("data/CV-Bench")) | |
| command.add_argument( | |
| "--selection", | |
| type=Path, | |
| default=Path("experiments/sites/zebra_rollout_viewer/scripts/cvbench_samples.json"), | |
| ) | |
| command.add_argument("--output-dir", type=Path, required=True) | |
| infer_parser = commands.choices["infer"] | |
| infer_parser.add_argument("--model-path", type=Path, required=True) | |
| infer_parser.add_argument( | |
| "--anole-inference-root", | |
| type=Path, | |
| default=Path("experiments/external/thinking-with-generated-images"), | |
| ) | |
| infer_parser.add_argument("--bagel-max-steps", type=int, default=12) | |
| infer_parser.add_argument("--seed", type=int, default=42) | |
| infer_parser.add_argument("--shard-id", type=int, default=0) | |
| infer_parser.add_argument("--num-shards", type=int, default=1) | |
| infer_parser.add_argument("--resume", action="store_true") | |
| return result | |
| def main() -> None: | |
| args = parser().parse_args() | |
| if args.command == "infer": | |
| if not 0 <= args.shard_id < args.num_shards: | |
| raise ValueError("shard-id must be in [0, num-shards)") | |
| infer(args) | |
| else: | |
| merge(args) | |
| if __name__ == "__main__": | |
| main() | |
Xet Storage Details
- Size:
- 7.7 kB
- Xet hash:
- 611e6c2bb4c9dfb1271151b65b8b85cb95f824a0699e7d37d9db0b2cfe63a5c3
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.