Buckets:
| """Run official Zebra-CoT models on disclosed-membership and paper-domain sets.""" | |
| from __future__ import annotations | |
| import argparse | |
| import io | |
| import json | |
| import re | |
| import time | |
| from pathlib import Path | |
| import pyarrow.parquet as pq | |
| from PIL import Image | |
| from . import anole, bagel | |
| from .common import read_jsonl, write_json | |
| from zebra_sft.scoring import score_answer | |
| from zebra_sft.trajectory import split_image_references | |
| PAPER_CONFIGS = ( | |
| "Scientific Reasoning - Graph Algorithms", | |
| "2D Visual Reasoning - Visual Jigsaw", | |
| "3D Visual Reasoning - Multi-Hop Objects Counting", | |
| "Visual Logic & Strategic Games - Maze", | |
| "Visual Logic & Strategic Games - Chess", | |
| "Visual Logic & Strategic Games - Tetris", | |
| ) | |
| def _manifest(path: Path) -> dict[str, dict]: | |
| return {item["sample_id"]: item for item in read_jsonl(path)} | |
| def _ours18_ids(root: Path) -> list[str]: | |
| rows = read_jsonl(root / "predictions.jsonl") | |
| if len(rows) != 18 or len({row["sample_id"] for row in rows}) != 18: | |
| raise ValueError("expected the canonical 18-sample grouped Zebra held-out run") | |
| return [row["sample_id"] for row in rows] | |
| def select_items(args: argparse.Namespace) -> list[dict]: | |
| manifest = _manifest(args.cache_root / "manifest.jsonl") | |
| if args.setting == "ours_grouped_test18": | |
| ids = _ours18_ids(args.ours18_root) | |
| missing = [sample_id for sample_id in ids if sample_id not in manifest] | |
| if missing: | |
| raise ValueError(f"ours18 IDs absent from cache manifest: {missing}") | |
| return [manifest[sample_id] for sample_id in ids] | |
| chosen = [] | |
| for config in PAPER_CONFIGS: | |
| candidates = [item for item in manifest.values() if item["config"] == config] | |
| if not candidates: | |
| raise ValueError(f"no cached examples for paper evaluation config {config}") | |
| chosen.append(min(candidates, key=lambda item: item["sample_id"])) | |
| return chosen | |
| def _read_parquet_row(path: Path, row_index: int) -> dict: | |
| parquet = pq.ParquetFile(path) | |
| offset = 0 | |
| for group in range(parquet.num_row_groups): | |
| count = parquet.metadata.row_group(group).num_rows | |
| if row_index < offset + count: | |
| return parquet.read_row_group(group).slice(row_index - offset, 1).to_pylist()[0] | |
| offset += count | |
| raise IndexError(f"row {row_index} is outside {path}") | |
| def _image(cell: dict, key: str) -> Image.Image: | |
| if not cell or not cell.get("bytes"): | |
| raise ValueError(f"missing observation image {key}") | |
| with Image.open(io.BytesIO(cell["bytes"])) as image: | |
| return image.convert("RGB").copy() | |
| def raw_item(item: dict, source_root: Path) -> tuple[str, list[Image.Image], str]: | |
| row = _read_parquet_row(source_root / item["source"]["shard"], item["source"]["row"]) | |
| parts = split_image_references(row["Question"]) | |
| observation_keys = [part["key"] for part in parts if part["type"] == "image"] | |
| question = "".join(part["text"] for part in parts if part["type"] == "text").strip() | |
| prompt = f"Question: {question}" | |
| return prompt, [_image(row[key], key) for key in observation_keys], row["Final Answer"] | |
| 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 _score(output: str, target: str) -> dict: | |
| # Bagel's native output wraps its explicit field in ``<answer>`` tags; | |
| # the UniAR-oriented shared scorer otherwise retains the closing tag. | |
| value = score_answer(re.sub(r"</?answer>", "", output, flags=re.IGNORECASE), target) | |
| value["raw_correct"] = value["correct"] | |
| value["has_final_answer"] = bool( | |
| re.search(r"final\s+answer\s*:|<answer>", output, flags=re.IGNORECASE) | |
| ) | |
| value["correct"] = bool(value["has_final_answer"] and value["raw_correct"]) | |
| return value | |
| def infer(args: argparse.Namespace) -> None: | |
| import random | |
| import numpy as np | |
| import torch | |
| items = select_items(args) | |
| selected = items[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["sample_id"] for row in read_jsonl(shard_path)} if shard_path.exists() and args.resume else set() | |
| pending = [item for item in selected if item["sample_id"] not in existing] | |
| runner, provenance = _runner(args) | |
| selection_meta = { | |
| "setting": args.setting, | |
| "sample_ids": [item["sample_id"] for item in items], | |
| "items": [{"sample_id": item["sample_id"], "config": item["config"], "source": item["source"]} for item in items], | |
| "meaning": ( | |
| "our deterministic grouped held-out split; these public source rows remain likely author-train members" | |
| if args.setting == "ours_grouped_test18" | |
| else "deterministic one-per-paper-evaluation-subtask proxy; authors did not publish the 512 IDs or their Anole in-distribution split" | |
| ), | |
| } | |
| write_json(args.output_dir / "selection.json", selection_meta) | |
| write_json( | |
| args.output_dir / f"shard_{args.shard_id:02d}_config.json", | |
| { | |
| "schema": "zebra_official_zebra_eval_v1", | |
| "model": args.model, | |
| "setting": args.setting, | |
| "selection": selection_meta, | |
| "model_path": str(args.model_path.resolve()), | |
| "provenance": provenance, | |
| }, | |
| ) | |
| with shard_path.open("a", encoding="utf-8") as handle: | |
| for position, item in enumerate(pending, start=1): | |
| sample_dir = args.output_dir / "samples" / item["sample_id"] | |
| sample_dir.mkdir(parents=True, exist_ok=True) | |
| sample_seed = args.seed + int(item["sample_id"][:8], 16) | |
| 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) | |
| prompt, images, target = raw_item(item, args.source_root) | |
| start = time.monotonic() | |
| output = ( | |
| runner.run(prompt, images, sample_dir, max_steps=args.bagel_max_steps) | |
| if args.model == "bagel" | |
| else runner.run(prompt, images, sample_dir) | |
| ) | |
| elapsed = time.monotonic() - start | |
| row = { | |
| "schema": "zebra_official_zebra_sample_v1", | |
| "model": args.model, | |
| "setting": args.setting, | |
| "sample_id": item["sample_id"], | |
| "config": item["config"], | |
| "source": item["source"], | |
| "prompt": prompt, | |
| "observation_count": len(images), | |
| "target": target, | |
| "seconds": elapsed, | |
| "seed": sample_seed, | |
| "provenance": provenance, | |
| **output, | |
| } | |
| row["score"] = _score(row["assistant_output"], target) | |
| (sample_dir / "output.txt").write_text(row["assistant_output"], encoding="utf-8") | |
| (sample_dir / "target.txt").write_text(target, 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), | |
| "sample_id": row["sample_id"], | |
| "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: | |
| items = select_items(args) | |
| expected = [item["sample_id"] for item in items] | |
| rows = {} | |
| for path in sorted((args.output_dir / "predictions").glob("shard_*.jsonl")): | |
| for row in read_jsonl(path): | |
| row["score"] = _score(row["assistant_output"], row["target"]) | |
| rows[row["sample_id"]] = row | |
| missing = [sample_id for sample_id in expected if sample_id not in rows] | |
| extra = sorted(set(rows) - set(expected)) | |
| if missing or extra: | |
| raise ValueError(f"incomplete Zebra setting: missing={missing}, extra={extra}") | |
| ordered = [rows[sample_id] for sample_id in expected] | |
| for row in ordered: | |
| write_json(args.output_dir / "samples" / row["sample_id"] / "result.json", row) | |
| (args.output_dir / "predictions.jsonl").write_text( | |
| "".join(json.dumps(row, ensure_ascii=False) + "\n" for row in ordered), encoding="utf-8" | |
| ) | |
| metrics = { | |
| "setting": args.setting, | |
| "model": args.model, | |
| "samples": len(ordered), | |
| "correct": sum(row["score"]["correct"] for row in ordered), | |
| "accuracy": sum(row["score"]["correct"] for row in ordered) / len(ordered), | |
| "generated_images": sum(len(row["generated_images"]) for row in ordered), | |
| "by_config": { | |
| row["config"]: {"correct": row["score"]["correct"], "generated_images": len(row["generated_images"])} | |
| for row in ordered | |
| }, | |
| } | |
| 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=("bagel", "anole"), required=True) | |
| command.add_argument("--setting", choices=("ours_grouped_test18", "paper_domain_proxy6"), required=True) | |
| command.add_argument("--output-dir", type=Path, required=True) | |
| command.add_argument("--source-root", type=Path, default=Path("data/Zebra-CoT")) | |
| command.add_argument("--cache-root", type=Path, default=Path("data/Zebra-CoT-UniAR-960")) | |
| command.add_argument( | |
| "--ours18-root", | |
| type=Path, | |
| default=Path("experiments/output/uniar-sft/zebra-cot-full/unversioned/runs/eval/20260821-111254_zebra_native_context_test18"), | |
| ) | |
| 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:
- 12.1 kB
- Xet hash:
- d7d043088410d6eda6562b0c4ec0f95d68259834e2199fd937749d7f507edf40
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.