Buckets:
| """Held-out Zebra evaluation: base direct answers versus fine-tuned dynamic rollout.""" | |
| from __future__ import annotations | |
| import argparse | |
| import json | |
| import math | |
| import re | |
| from collections import defaultdict | |
| from pathlib import Path | |
| import numpy as np | |
| import torch | |
| from peft import PeftModel | |
| from PIL import Image, ImageDraw | |
| from torch.utils.tensorboard import SummaryWriter | |
| from transformers import AutoProcessor | |
| from artifacts import is_timestamped_run_name | |
| from uniar import UniARForConditionalGeneration, UniARVisualDecoder, configure_processor | |
| from .cached_rollout import CachedInterleavedRollout | |
| from .dataset import image_gen_token_id | |
| from .full_data import cached_batch, load_item_images | |
| from .generated_codes import save_generated_codes | |
| from .objective import joint_sft_loss | |
| from .scoring import score_answer | |
| def _write_json(path: Path, value) -> None: | |
| path.write_text( | |
| json.dumps(value, indent=2, ensure_ascii=False, default=str) + "\n", | |
| encoding="utf-8", | |
| ) | |
| def _load_shard(path: Path) -> list[dict]: | |
| return torch.load(path, map_location="cpu", weights_only=False, mmap=True)["items"] | |
| def _select_items(cache_root: Path, split: str, count: int) -> list[dict]: | |
| rows = [] | |
| with (cache_root / "manifest.jsonl").open(encoding="utf-8") as handle: | |
| for line in handle: | |
| row = json.loads(line) | |
| if row["split"] == split: | |
| rows.append(row) | |
| by_config: dict[str, list[dict]] = defaultdict(list) | |
| for row in rows: | |
| by_config[row["config"]].append(row) | |
| selected = [] | |
| # Prefer inexpensive full-trajectory examples while covering all 18 configs. | |
| for config in sorted(by_config): | |
| selected.append( | |
| min( | |
| by_config[config], | |
| key=lambda row: ( | |
| row["generated_images"], | |
| row["visual_positions"], | |
| row["sequence_length"], | |
| row["split_score"], | |
| ), | |
| ) | |
| ) | |
| selected_ids = {row["sample_id"] for row in selected} | |
| for row in sorted( | |
| rows, | |
| key=lambda row: ( | |
| row["generated_images"], | |
| row["visual_positions"], | |
| row["sequence_length"], | |
| row["split_score"], | |
| ), | |
| ): | |
| if len(selected) >= count: | |
| break | |
| if row["sample_id"] not in selected_ids: | |
| selected.append(row) | |
| selected_ids.add(row["sample_id"]) | |
| cache = {} | |
| items = [] | |
| for row in selected[:count]: | |
| if row["cache_shard"] not in cache: | |
| cache[row["cache_shard"]] = _load_shard(cache_root / row["cache_shard"]) | |
| items.append(cache[row["cache_shard"]][row["cache_index"]]) | |
| return items | |
| def _target(item: dict) -> str: | |
| return next( | |
| event["text"] | |
| for event in item["trajectory"]["assistant"] | |
| if event["role"] == "answer" | |
| ) | |
| def _teacher_forced(model, processor, items: list[dict], args) -> dict: | |
| marker = image_gen_token_id(processor) | |
| rows = [] | |
| for item in items: | |
| batch, codes = cached_batch(item) | |
| with torch.autocast("cuda", dtype=torch.bfloat16): | |
| result = joint_sft_loss( | |
| model, | |
| batch, | |
| codes, | |
| image_gen_id=marker, | |
| visual_loss_weight=args.visual_loss_weight, | |
| marker_token_weight=args.marker_token_weight, | |
| ) | |
| rows.append( | |
| { | |
| key: float(result[key].detach().float().cpu()) | |
| for key in ( | |
| "loss", | |
| "text_loss", | |
| "marker_loss", | |
| "visual_loss", | |
| "text_accuracy", | |
| "marker_accuracy", | |
| "visual_accuracy", | |
| ) | |
| } | |
| | {"sample_id": item["sample_id"], "config": item["config"]} | |
| ) | |
| aggregate = { | |
| key: sum(row[key] for row in rows if not math.isnan(row[key])) | |
| / sum(not math.isnan(row[key]) for row in rows) | |
| for key in rows[0] | |
| if key not in {"sample_id", "config"} | |
| } | |
| return aggregate | {"samples": len(rows), "rows": rows} | |
| def _summarize(rows: list[dict]) -> dict: | |
| by_config = defaultdict(lambda: {"samples": 0, "correct": 0}) | |
| for row in rows: | |
| entry = by_config[row["config"]] | |
| entry["samples"] += 1 | |
| entry["correct"] += int(row["score"]["correct"]) | |
| for entry in by_config.values(): | |
| entry["accuracy"] = entry["correct"] / entry["samples"] | |
| return { | |
| "samples": len(rows), | |
| "correct": sum(row["score"]["correct"] for row in rows), | |
| "accuracy": sum(row["score"]["correct"] for row in rows) / len(rows), | |
| "by_config": dict(sorted(by_config.items())), | |
| } | |
| def _readable(text: str) -> str: | |
| counter = 0 | |
| def replace(match): | |
| nonlocal counter | |
| counter += 1 | |
| return f"<generated_visual_{counter:02d}:{match.group(0).count('<|image_pad|>')}_positions>" | |
| return re.sub(r"(?:<\|image_pad\|>)+", replace, text) | |
| def _contact(target: Image.Image, generated: Image.Image, title: str) -> Image.Image: | |
| size = 384 | |
| canvas = Image.new("RGB", (size * 2, size + 32), "white") | |
| canvas.paste(target.convert("RGB").resize((size, size)), (0, 32)) | |
| canvas.paste(generated.convert("RGB").resize((size, size)), (size, 32)) | |
| draw = ImageDraw.Draw(canvas) | |
| draw.text((4, 8), f"GT state: {title}", fill="black") | |
| draw.text((size + 4, 8), "autonomous generated state", fill="black") | |
| return canvas | |
| def evaluate(args: argparse.Namespace) -> dict: | |
| if not is_timestamped_run_name(args.output_dir.name): | |
| raise ValueError("output directory must use a canonical timestamped run name") | |
| args.output_dir.mkdir(parents=True, exist_ok=False) | |
| items = _select_items(args.cache_root, args.split, args.samples) | |
| processor = configure_processor( | |
| AutoProcessor.from_pretrained(args.model_path, use_fast=True) | |
| ) | |
| base = UniARForConditionalGeneration.from_pretrained( | |
| args.model_path, | |
| torch_dtype=torch.bfloat16, | |
| attn_implementation=args.attn, | |
| low_cpu_mem_usage=True, | |
| ).to("cuda").eval() | |
| base_teacher = _teacher_forced(base, processor, items, args) | |
| base_rollout = CachedInterleavedRollout(base, processor) | |
| baseline_rows = [] | |
| for index, item in enumerate(items): | |
| result = base_rollout.direct_answer(item, max_new_tokens=args.max_text_tokens) | |
| score = score_answer(result["decoded"], _target(item)) | |
| row = { | |
| "sample_id": item["sample_id"], | |
| "config": item["config"], | |
| "target": _target(item), | |
| "decoded": result["decoded"], | |
| "sequence_ids": result["sequence_ids"].tolist(), | |
| "stop_reason": result["stop_reason"], | |
| "score": score, | |
| } | |
| baseline_rows.append(row) | |
| print( | |
| f"BASELINE {index + 1}/{len(items)} config={item['config']} " | |
| f"correct={score['correct']}", | |
| flush=True, | |
| ) | |
| _write_json(args.output_dir / "base_teacher_forced.json", base_teacher) | |
| with (args.output_dir / "baseline_predictions.jsonl").open("w", encoding="utf-8") as handle: | |
| for row in baseline_rows: | |
| handle.write(json.dumps(row, ensure_ascii=False) + "\n") | |
| model = PeftModel.from_pretrained(base, args.adapter_dir).eval() | |
| finetuned_teacher = _teacher_forced(model, processor, items, args) | |
| rollout = CachedInterleavedRollout(model, processor) | |
| fine_rows = [] | |
| raw_outputs = {} | |
| for index, item in enumerate(items): | |
| result = rollout.run( | |
| item, | |
| max_text_tokens=args.max_text_tokens, | |
| max_scaffold_tokens=args.max_scaffold_tokens, | |
| ) | |
| raw_codes = result.pop("raw_codes") | |
| target_codes = result.pop("target_codes") | |
| counts = item["generated_image_token_counts"].tolist() | |
| target_steps = torch.split(target_codes, counts) | |
| visual_steps = [] | |
| for step, codes in enumerate(raw_codes): | |
| generated_grid = result["generated_grids"][step] | |
| expected_grid = result["expected_grids"][step] if step < len(result["expected_grids"]) else None | |
| bit_accuracy = None | |
| if expected_grid == generated_grid and step < len(target_steps): | |
| bit_accuracy = float( | |
| (codes[0].to(torch.uint8) == target_steps[step]).float().mean() | |
| ) | |
| visual_steps.append( | |
| { | |
| "step": step + 1, | |
| "expected_grid": expected_grid, | |
| "generated_grid": generated_grid, | |
| "bit_accuracy": bit_accuracy, | |
| } | |
| ) | |
| score = score_answer(result["decoded"], _target(item)) | |
| protocol_ok = ( | |
| result["stop_reason"] == "im_end" | |
| and result["generated_grids"] == result["expected_grids"] | |
| ) | |
| row = { | |
| **{key: value for key, value in result.items() if key != "sequence_ids"}, | |
| "sequence_ids": result["sequence_ids"].tolist(), | |
| "readable_sequence": _readable(result["decoded"]), | |
| "config": item["config"], | |
| "target": _target(item), | |
| "protocol_ok": protocol_ok, | |
| "score": score, | |
| "visual_steps": visual_steps, | |
| } | |
| fine_rows.append(row) | |
| raw_outputs[item["sample_id"]] = raw_codes | |
| sample_dir = args.output_dir / "samples" / item["sample_id"] | |
| sample_dir.mkdir(parents=True) | |
| (sample_dir / "target.txt").write_text(item["target_render"], encoding="utf-8") | |
| (sample_dir / "output.txt").write_text(row["readable_sequence"], encoding="utf-8") | |
| save_generated_codes(sample_dir / "generated_codes.pt", raw_codes) | |
| _write_json(sample_dir / "result.json", row) | |
| print( | |
| f"PROGRESS {index + 1}/{len(items)} config={item['config']} " | |
| f"correct={score['correct']} protocol={protocol_ok} " | |
| f"images={result['generated_images']}/{result['expected_images']}", | |
| flush=True, | |
| ) | |
| _write_json(args.output_dir / "finetuned_teacher_forced.json", finetuned_teacher) | |
| with (args.output_dir / "finetuned_predictions.jsonl").open("w", encoding="utf-8") as handle: | |
| for row in fine_rows: | |
| handle.write(json.dumps(row, ensure_ascii=False) + "\n") | |
| del rollout, base_rollout, model, base | |
| torch.cuda.empty_cache() | |
| decoder = UniARVisualDecoder.from_pretrained(args.model_path, device=torch.device("cuda")) | |
| writer = SummaryWriter(args.training_run / "tensorboard") | |
| training_summary = json.loads((args.training_run / "summary.json").read_text()) | |
| tensorboard_step = training_summary["optimizer_steps"] | |
| galleries = 0 | |
| for item, row in zip(items, fine_rows): | |
| writer.add_text( | |
| f"heldout_outputs/{item['config'][:48]}_{item['sample_id']}", | |
| "```text\n" + row["readable_sequence"] + "\n```", | |
| tensorboard_step, | |
| ) | |
| raw_codes = raw_outputs[item["sample_id"]] | |
| if galleries >= args.gallery_samples or not raw_codes: | |
| continue | |
| codes = raw_codes[0].to("cuda", torch.bfloat16) | |
| indices = decoder.bsq_encoder.bsq.codes_to_indexes(codes * 2 - 1) | |
| grid_h, grid_w = row["generated_grids"][0] | |
| generated = decoder.decode( | |
| indices, | |
| ar_height=grid_h * 32, | |
| ar_width=grid_w * 32, | |
| upsampling_ratio=1.0, | |
| num_inference_steps=args.decoder_steps, | |
| guidance_scale=args.decoder_guidance, | |
| )[0] | |
| image_key = next( | |
| event["key"] | |
| for event in item["trajectory"]["assistant"] | |
| if event["role"] == "generated_state" | |
| ) | |
| target = load_item_images(item, args.source_root, args.graph_root)[image_key] | |
| contact = _contact(target, generated, item["config"]) | |
| sample_dir = args.output_dir / "samples" / item["sample_id"] | |
| generated.save(sample_dir / "generated_01.png") | |
| contact.save(sample_dir / "contact_sheet.png") | |
| writer.add_image( | |
| f"heldout_images/{item['config'][:48]}_{item['sample_id']}", | |
| np.asarray(contact), | |
| tensorboard_step, | |
| dataformats="HWC", | |
| ) | |
| galleries += 1 | |
| writer.flush() | |
| writer.close() | |
| baseline = _summarize(baseline_rows) | |
| finetuned = _summarize(fine_rows) | |
| summary = { | |
| "split": args.split, | |
| "selection": "one low-compute full trajectory per config, then globally lowest compute", | |
| "base_teacher_forced": {key: value for key, value in base_teacher.items() if key != "rows"}, | |
| "finetuned_teacher_forced": { | |
| key: value for key, value in finetuned_teacher.items() if key != "rows" | |
| }, | |
| "baseline_direct": baseline, | |
| "finetuned_interleaved": finetuned, | |
| "answer_accuracy_delta": finetuned["accuracy"] - baseline["accuracy"], | |
| "protocol_exact": sum(row["protocol_ok"] for row in fine_rows), | |
| "protocol_accuracy": sum(row["protocol_ok"] for row in fine_rows) / len(fine_rows), | |
| "generated_images": sum(row["generated_images"] for row in fine_rows), | |
| "expected_images": sum(row["expected_images"] for row in fine_rows), | |
| "tensorboard": str(args.training_run / "tensorboard"), | |
| "gallery_samples": galleries, | |
| } | |
| _write_json(args.output_dir / "summary.json", summary) | |
| print("EVALUATION_COMPLETE " + json.dumps(summary), flush=True) | |
| return summary | |
| def build_parser() -> argparse.ArgumentParser: | |
| parser = argparse.ArgumentParser(description=__doc__) | |
| parser.add_argument( | |
| "--model-path", type=Path, default=Path("models/UniAR-SFT") | |
| ) | |
| parser.add_argument("--adapter-dir", type=Path, required=True) | |
| parser.add_argument("--training-run", type=Path, required=True) | |
| parser.add_argument("--cache-root", type=Path, default=Path("data/Zebra-CoT-UniAR-960")) | |
| parser.add_argument("--source-root", type=Path, default=Path("data/Zebra-CoT")) | |
| parser.add_argument( | |
| "--graph-root", type=Path, default=Path("/mnt/ultracube/datasets/Graph-Algorithms") | |
| ) | |
| parser.add_argument("--output-dir", type=Path, required=True) | |
| parser.add_argument("--split", default="test", choices=["validation", "test"]) | |
| parser.add_argument("--samples", type=int, default=18) | |
| parser.add_argument("--gallery-samples", type=int, default=6) | |
| parser.add_argument("--max-text-tokens", type=int, default=512) | |
| parser.add_argument("--max-scaffold-tokens", type=int, default=16) | |
| parser.add_argument("--decoder-steps", type=int, default=28) | |
| parser.add_argument("--decoder-guidance", type=float, default=1.5) | |
| parser.add_argument("--visual-loss-weight", type=float, default=10.0) | |
| parser.add_argument("--marker-token-weight", type=float, default=5.0) | |
| parser.add_argument("--attn", default="flash_attention_2") | |
| return parser | |
| def main() -> None: | |
| evaluate(build_parser().parse_args()) | |
| if __name__ == "__main__": | |
| main() | |
Xet Storage Details
- Size:
- 15.3 kB
- Xet hash:
- 412f9addad801dd7bb4eaa87e124284fca3b92ebdd27ec89a091d0b080a3b1ce
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.