Buckets:
| """Run causal visual-context ablations on the fixed 18-sample Zebra test audit.""" | |
| from __future__ import annotations | |
| import argparse | |
| import gc | |
| import hashlib | |
| import json | |
| from collections import Counter | |
| from pathlib import Path | |
| import torch | |
| from peft import PeftModel | |
| from transformers import AutoProcessor | |
| from artifacts import is_timestamped_run_name | |
| from uniar import UniARForConditionalGeneration, configure_processor | |
| from .cached_rollout import ABLATIONS, CachedInterleavedRollout | |
| from .evaluate_corpus import _readable, _target, _write_json | |
| from .generated_codes import save_generated_codes | |
| from .retry_rollouts import _load_items | |
| from .scoring import score_answer | |
| CONDITIONS = tuple(sorted(ABLATIONS - {"control"})) | |
| def _selection_rows(root: Path) -> list[dict]: | |
| path = root / "predictions.jsonl" | |
| rows = [json.loads(line) for line in path.read_text().splitlines() if line] | |
| if len(rows) != 18 or len({row["sample_id"] for row in rows}) != 18: | |
| raise ValueError(f"expected exactly 18 unique fixed samples in {path}") | |
| return rows | |
| def _code_hash(codes: torch.Tensor | None) -> str | None: | |
| if codes is None: | |
| return None | |
| return hashlib.sha256(codes.contiguous().numpy().tobytes()).hexdigest() | |
| def _sample_complete(sample_dir: Path, condition: str) -> bool: | |
| required = ["result.json", "target.txt", "output.txt", "generated_codes.pt"] | |
| if not all((sample_dir / name).is_file() for name in required): | |
| return False | |
| row = json.loads((sample_dir / "result.json").read_text()) | |
| if row.get("ablation") != condition: | |
| return False | |
| if condition in {"noise", "gray"}: | |
| return all( | |
| (sample_dir / f"injected_{step:02d}.png").is_file() | |
| for step in range(1, row["generated_images"] + 1) | |
| ) | |
| return True | |
| def _visual_steps(item: dict, result: dict, raw_codes: list) -> list[dict]: | |
| target_codes = result.pop("target_codes") | |
| target_steps = torch.split(target_codes, item["generated_image_token_counts"].tolist()) | |
| rows = [] | |
| for step, (grid, codes) in enumerate( | |
| zip(result["generated_grids"], raw_codes), start=1 | |
| ): | |
| expected = ( | |
| result["expected_grids"][step - 1] | |
| if step <= len(result["expected_grids"]) | |
| else None | |
| ) | |
| bit_accuracy = None | |
| if codes is not None and expected == grid and step <= len(target_steps): | |
| bit_accuracy = float( | |
| (codes[0].to(torch.uint8) == target_steps[step - 1]).float().mean() | |
| ) | |
| rows.append( | |
| { | |
| "step": step, | |
| "declared_grid": grid, | |
| "expected_grid": expected, | |
| "serialized_positions": next( | |
| event["serialized_positions"] | |
| for event in result["events"] | |
| if event["type"] == "image" and event["index"] == step | |
| ), | |
| "bit_accuracy": bit_accuracy, | |
| "code_sha256": _code_hash(codes), | |
| } | |
| ) | |
| return rows | |
| def run(args: argparse.Namespace) -> None: | |
| if not is_timestamped_run_name(args.suite_dir.name): | |
| raise ValueError("suite directory must use a canonical timestamped run name") | |
| rows = _selection_rows(args.selection_root) | |
| selected_ids = [row["sample_id"] for row in rows] | |
| if args.sample_id: | |
| requested = set(args.sample_id) | |
| missing = requested - set(selected_ids) | |
| if missing: | |
| raise ValueError(f"unknown sample IDs: {sorted(missing)}") | |
| selected_ids = [sample_id for sample_id in selected_ids if sample_id in requested] | |
| selected_ids = selected_ids[args.shard_index :: args.shard_count] | |
| items = _load_items(args.cache_root, selected_ids) | |
| 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() | |
| model = PeftModel.from_pretrained(base, args.adapter_dir).eval() | |
| rollout = CachedInterleavedRollout(model, processor) | |
| totals = {} | |
| for condition in args.condition: | |
| output_root = args.suite_dir / condition | |
| sample_root = output_root / "samples" | |
| job_root = output_root / "jobs" | |
| sample_root.mkdir(parents=True, exist_ok=True) | |
| job_root.mkdir(parents=True, exist_ok=True) | |
| job_path = job_root / f"shard_{args.shard_index:02d}-of-{args.shard_count:02d}.json" | |
| if job_path.exists(): | |
| if not args.resume: | |
| raise FileExistsError(job_path) | |
| manifest = json.loads(job_path.read_text()) | |
| if manifest.get("sample_ids") != [item["sample_id"] for item in items]: | |
| raise ValueError(f"existing job manifest does not match selection: {job_path}") | |
| print( | |
| f"ABLATION_CONDITION_ALREADY_COMPLETE condition={condition} path={job_path}", | |
| flush=True, | |
| ) | |
| totals[condition] = len(items) | |
| continue | |
| completed = [] | |
| for index, item in enumerate(items, start=1): | |
| sample_dir = sample_root / item["sample_id"] | |
| result_path = sample_dir / "result.json" | |
| if args.resume and _sample_complete(sample_dir, condition): | |
| completed.append(item["sample_id"]) | |
| print( | |
| f"PROGRESS {index}/{len(items)} condition={condition} " | |
| f"sample={item['sample_id']} status=already_complete", | |
| flush=True, | |
| ) | |
| continue | |
| if sample_dir.exists() and any(sample_dir.iterdir()): | |
| raise FileExistsError( | |
| f"refusing to reuse incomplete non-empty sample directory: {sample_dir}" | |
| ) | |
| sample_dir.mkdir(parents=True, exist_ok=True) | |
| def report_image(event: dict) -> None: | |
| print( | |
| f"IMAGE_PROGRESS condition={condition} sample={item['sample_id']} " | |
| f"step={event['index']} grid={event['grid'][0]}x{event['grid'][1]} " | |
| f"context={event['sequence_length']}/{rollout.context_limit}", | |
| flush=True, | |
| ) | |
| result = rollout.run( | |
| item, | |
| max_text_tokens=None, | |
| max_scaffold_tokens=None, | |
| extra_image_budget=None, | |
| cycle_repeats=args.cycle_repeats, | |
| ablation=condition, | |
| replacement_seed=args.seed, | |
| event_callback=report_image, | |
| ) | |
| raw_codes = result.pop("raw_codes") | |
| intervention_images = result.pop("intervention_images") | |
| visual_steps = _visual_steps(item, result, raw_codes) | |
| decoded = result["decoded"] | |
| scaffolds = [event for event in result["events"] if event["type"] == "scaffold"] | |
| scaffold_ok = all( | |
| event.get("event") == "vision_start" and event.get("grid") is not None | |
| for event in scaffolds | |
| ) | |
| has_answer = "Final answer:" in decoded | |
| relaxed_complete = result["stop_reason"] == "im_end" and has_answer and scaffold_ok | |
| score = score_answer(decoded, _target(item)) | |
| row = { | |
| **{key: value for key, value in result.items() if key != "sequence_ids"}, | |
| "sequence_ids": result["sequence_ids"].tolist(), | |
| "readable_sequence": _readable(decoded), | |
| "config": item["config"], | |
| "target": _target(item), | |
| "scaffold_ok": scaffold_ok, | |
| "has_final_answer": has_answer, | |
| "balanced_think_tags": decoded.count("<think>") == decoded.count("</think>"), | |
| "relaxed_complete": relaxed_complete, | |
| "score": score, | |
| "visual_steps": visual_steps, | |
| "replacement_seed": args.seed, | |
| } | |
| (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) | |
| for step, image in enumerate(intervention_images, start=1): | |
| if image is not None: | |
| image.save(sample_dir / f"injected_{step:02d}.png") | |
| # result.json is the completion marker and is intentionally written | |
| # only after every companion artifact is durable. | |
| _write_json(sample_dir / "result.json", row) | |
| completed.append(item["sample_id"]) | |
| print( | |
| f"PROGRESS {index}/{len(items)} condition={condition} " | |
| f"config={item['config']} stop={result['stop_reason']} " | |
| f"correct={score['correct']} images={result['generated_images']}", | |
| flush=True, | |
| ) | |
| del result, raw_codes, intervention_images | |
| gc.collect() | |
| torch.cuda.empty_cache() | |
| _write_json( | |
| job_path, | |
| { | |
| "condition": condition, | |
| "context_limit": rollout.context_limit, | |
| "cycle_repeats": args.cycle_repeats, | |
| "seed": args.seed, | |
| "shard_index": args.shard_index, | |
| "shard_count": args.shard_count, | |
| "sample_ids": completed, | |
| }, | |
| ) | |
| totals[condition] = len(completed) | |
| print("ABLATION_JOB_COMPLETE " + json.dumps(totals), flush=True) | |
| def aggregate(args: argparse.Namespace) -> None: | |
| selection = _selection_rows(args.selection_root) | |
| expected = {row["sample_id"] for row in selection} | |
| summaries = {} | |
| for condition in args.condition or CONDITIONS: | |
| paths = sorted((args.suite_dir / condition / "samples").glob("*/result.json")) | |
| rows = [json.loads(path.read_text()) for path in paths] | |
| actual = {row["sample_id"] for row in rows} | |
| if actual != expected: | |
| raise ValueError( | |
| f"{condition}: missing={sorted(expected - actual)}, extra={sorted(actual - expected)}" | |
| ) | |
| if any(row.get("ablation") != condition for row in rows): | |
| raise ValueError(f"{condition}: result condition mismatch") | |
| if {row.get("context_limit") for row in rows} != {32768}: | |
| raise ValueError(f"{condition}: expected every context_limit to be 32768") | |
| summaries[condition] = { | |
| "samples": len(rows), | |
| "correct": sum( | |
| bool(row.get("has_final_answer") and row["score"]["correct"]) | |
| for row in rows | |
| ), | |
| "relaxed_complete": sum(row["relaxed_complete"] for row in rows), | |
| "generated_images": sum(row["generated_images"] for row in rows), | |
| "mean_assistant_tokens": sum(len(row["sequence_ids"]) for row in rows) / len(rows), | |
| "stop_reasons": dict(sorted(Counter(row["stop_reason"] for row in rows).items())), | |
| } | |
| _write_json(args.suite_dir / condition / "summary.json", summaries[condition]) | |
| control_rows = _selection_rows(args.control_root) | |
| if [row["sample_id"] for row in control_rows] != [ | |
| row["sample_id"] for row in selection | |
| ]: | |
| raise ValueError("control and selection sample order differs") | |
| suite = { | |
| "schema": "zebra_visual_context_ablation_v1", | |
| "split": "held-out test", | |
| "model_path": str(args.model_path.resolve()), | |
| "adapter_dir": str(args.adapter_dir.resolve()), | |
| "selection_root": str(args.selection_root.resolve()), | |
| "control_root": str(args.control_root.resolve()), | |
| "sample_ids": [row["sample_id"] for row in selection], | |
| "context_limit": 32768, | |
| "generation_limits": { | |
| "text_tokens_per_segment": None, | |
| "scaffold_tokens_per_segment": None, | |
| "generated_images": None, | |
| "cycle_repeats": 3, | |
| "cycle_max_period": 4, | |
| }, | |
| "control": { | |
| "samples": len(control_rows), | |
| "correct": sum( | |
| bool(row.get("has_final_answer") and row["score"]["correct"]) | |
| for row in control_rows | |
| ), | |
| "relaxed_complete": sum(row["relaxed_complete"] for row in control_rows), | |
| "generated_images": sum(row["generated_images"] for row in control_rows), | |
| "stop_reasons": dict( | |
| sorted(Counter(row["stop_reason"] for row in control_rows).items()) | |
| ), | |
| }, | |
| "conditions": summaries, | |
| } | |
| _write_json(args.suite_dir / "summary.json", suite) | |
| print("ABLATION_SUITE_COMPLETE " + json.dumps(suite, ensure_ascii=False), flush=True) | |
| def build_parser() -> argparse.ArgumentParser: | |
| parser = argparse.ArgumentParser(description=__doc__) | |
| subparsers = parser.add_subparsers(dest="command", required=True) | |
| run_parser = subparsers.add_parser("run") | |
| run_parser.add_argument( | |
| "--condition", action="append", required=True, choices=tuple(sorted(ABLATIONS)) | |
| ) | |
| run_parser.add_argument("--suite-dir", type=Path, required=True) | |
| run_parser.add_argument("--selection-root", type=Path, required=True) | |
| run_parser.add_argument("--cache-root", type=Path, required=True) | |
| run_parser.add_argument("--adapter-dir", type=Path, required=True) | |
| run_parser.add_argument( | |
| "--model-path", type=Path, default=Path("models/UniAR-SFT") | |
| ) | |
| run_parser.add_argument("--sample-id", action="append") | |
| run_parser.add_argument("--shard-index", type=int, default=0) | |
| run_parser.add_argument("--shard-count", type=int, default=1) | |
| run_parser.add_argument("--cycle-repeats", type=int, default=3) | |
| run_parser.add_argument("--seed", type=int, default=20260823) | |
| run_parser.add_argument("--attn", default="flash_attention_2") | |
| run_parser.add_argument( | |
| "--resume", | |
| action="store_true", | |
| help="skip complete samples and reuse only empty interrupted sample directories", | |
| ) | |
| aggregate_parser = subparsers.add_parser("aggregate") | |
| aggregate_parser.add_argument("--suite-dir", type=Path, required=True) | |
| aggregate_parser.add_argument("--selection-root", type=Path, required=True) | |
| aggregate_parser.add_argument("--control-root", type=Path, required=True) | |
| aggregate_parser.add_argument("--adapter-dir", type=Path, required=True) | |
| aggregate_parser.add_argument( | |
| "--model-path", type=Path, default=Path("models/UniAR-SFT") | |
| ) | |
| aggregate_parser.add_argument("--condition", action="append", choices=CONDITIONS) | |
| return parser | |
| def main() -> None: | |
| args = build_parser().parse_args() | |
| if args.command == "run": | |
| if not 0 <= args.shard_index < args.shard_count: | |
| raise ValueError("shard-index must be in [0, shard-count)") | |
| run(args) | |
| else: | |
| aggregate(args) | |
| if __name__ == "__main__": | |
| main() | |
Xet Storage Details
- Size:
- 15.2 kB
- Xet hash:
- 94c96677d8790fc6dc2bc6575089d9896559a6c1be6ebda458daa19c6643526e
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.