Buckets:
| """Create the inspectable 16-example Zebra-CoT overfit fixture.""" | |
| from __future__ import annotations | |
| import argparse | |
| import hashlib | |
| import io | |
| import json | |
| import shutil | |
| from pathlib import Path | |
| from .trajectory import IMAGE_REFERENCE, build_trajectory, symbolic_render | |
| CONFIG = "3D Visual Reasoning - Multi-Hop Objects Counting" | |
| def _write_json(path: Path, value) -> None: | |
| path.write_text(json.dumps(value, indent=2, ensure_ascii=False) + "\n", encoding="utf-8") | |
| def _letterbox(image, size: int): | |
| from PIL import Image | |
| image = image.convert("RGB") | |
| scale = min(size / image.width, size / image.height) | |
| resized = image.resize( | |
| (max(1, round(image.width * scale)), max(1, round(image.height * scale))), | |
| Image.Resampling.LANCZOS, | |
| ) | |
| canvas = Image.new("RGB", (size, size), (127, 127, 127)) | |
| canvas.paste(resized, ((size - resized.width) // 2, (size - resized.height) // 2)) | |
| return canvas | |
| def prepare_samples( | |
| source_root: Path, | |
| output_dir: Path, | |
| *, | |
| count: int = 16, | |
| image_size: int = 256, | |
| overwrite: bool = False, | |
| ) -> dict: | |
| """Copy exact raw rows and emit normalized typed trajectories.""" | |
| import pyarrow.parquet as pq | |
| from PIL import Image | |
| if image_size % 32: | |
| raise ValueError("image_size must be divisible by UniAR's downsample factor 32") | |
| if output_dir.exists(): | |
| if not overwrite: | |
| raise FileExistsError(f"refusing to overwrite {output_dir}; pass --overwrite") | |
| shutil.rmtree(output_dir) | |
| source_dir = source_root / CONFIG | |
| shard = sorted(source_dir.glob("train-*.parquet"))[0] | |
| table = pq.read_table(shard).slice(0, count) | |
| if table.num_rows != count: | |
| raise ValueError(f"requested {count} rows but found {table.num_rows}") | |
| raw_dir = output_dir / "raw" | |
| prepared_dir = output_dir / "prepared" | |
| images_dir = prepared_dir / "images" | |
| rendered_dir = prepared_dir / "rendered" | |
| raw_dir.mkdir(parents=True) | |
| images_dir.mkdir(parents=True) | |
| rendered_dir.mkdir(parents=True) | |
| pq.write_table(table, raw_dir / "samples.parquet", compression="zstd") | |
| trajectories: list[dict] = [] | |
| source_records: list[dict] = [] | |
| image_count = 0 | |
| for row_index, row in enumerate(table.to_pylist()): | |
| sample_id = f"zebra-3d-counting-{row_index:06d}" | |
| refs = IMAGE_REFERENCE.findall(row["Question"] + "\n" + row["Text Reasoning Trace"]) | |
| if len(refs) != len(set(refs)): | |
| raise ValueError(f"duplicate image reference in {sample_id}") | |
| image_paths: dict[str, str] = {} | |
| image_meta: dict[str, dict] = {} | |
| sample_image_dir = images_dir / sample_id | |
| sample_image_dir.mkdir() | |
| for key in refs: | |
| value = row.get(key) | |
| if not value or not value.get("bytes"): | |
| raise ValueError(f"missing {key} in {sample_id}") | |
| raw_bytes = value["bytes"] | |
| with Image.open(io.BytesIO(raw_bytes)) as image: | |
| original = {"width": image.width, "height": image.height, "mode": image.mode} | |
| normalized = _letterbox(image, image_size) | |
| relative = Path("prepared") / "images" / sample_id / f"{key}.png" | |
| normalized.save(output_dir / relative, format="PNG", optimize=True) | |
| image_paths[key] = relative.as_posix() | |
| image_meta[key] = { | |
| **original, | |
| "sha256": hashlib.sha256(raw_bytes).hexdigest(), | |
| "prepared_path": relative.as_posix(), | |
| } | |
| image_count += 1 | |
| trajectory = build_trajectory( | |
| sample_id=sample_id, | |
| config=CONFIG, | |
| source_shard=shard.name, | |
| source_row=row_index, | |
| question=row["Question"], | |
| reasoning_trace=row["Text Reasoning Trace"], | |
| final_answer=row["Final Answer"], | |
| image_paths=image_paths, | |
| image_size=image_size, | |
| ) | |
| trajectories.append(trajectory) | |
| source_records.append( | |
| { | |
| "sample_id": sample_id, | |
| "source_shard": shard.name, | |
| "source_row": row_index, | |
| "Question": row["Question"], | |
| "Text Reasoning Trace": row["Text Reasoning Trace"], | |
| "Final Answer": row["Final Answer"], | |
| "images": image_meta, | |
| } | |
| ) | |
| (rendered_dir / f"{sample_id}.txt").write_text( | |
| symbolic_render(trajectory), encoding="utf-8" | |
| ) | |
| with (raw_dir / "records.jsonl").open("w", encoding="utf-8") as handle: | |
| for record in source_records: | |
| handle.write(json.dumps(record, ensure_ascii=False) + "\n") | |
| with (prepared_dir / "trajectories.jsonl").open("w", encoding="utf-8") as handle: | |
| for trajectory in trajectories: | |
| handle.write(json.dumps(trajectory, ensure_ascii=False) + "\n") | |
| manifest = { | |
| "dataset": "Zebra-CoT", | |
| "config": CONFIG, | |
| "source_root": str(source_root.resolve()), | |
| "source_shard": shard.name, | |
| "source_rows": [0, count - 1], | |
| "samples": count, | |
| "images": image_count, | |
| "image_transform": { | |
| "method": "aspect-preserving letterbox", | |
| "size": [image_size, image_size], | |
| "padding_rgb": [127, 127, 127], | |
| }, | |
| "raw_format": "raw/samples.parquet preserves the source row schema and embedded images", | |
| "prepared_format": "prepared/trajectories.jsonl uses spatial_mmcot_trajectory_v1", | |
| } | |
| _write_json(output_dir / "manifest.json", manifest) | |
| (output_dir / "README.md").write_text( | |
| "# Zebra-CoT Samples\n\n" | |
| "Deterministic 16-trajectory overfit fixture from `3D Visual Reasoning - " | |
| "Multi-Hop Objects Counting`.\n\n" | |
| "- `raw/samples.parquet`: exact source-format rows with embedded images.\n" | |
| "- `raw/records.jsonl`: inspectable original text plus image hashes.\n" | |
| "- `prepared/trajectories.jsonl`: typed observation/thought/generated-state/answer events.\n" | |
| "- `prepared/images/`: 256x256 aspect-preserving normalized PNGs.\n" | |
| "- `prepared/rendered/`: human-readable UniAR sequences with `<image_gen>` boundaries.\n\n" | |
| "The source Parquet is never modified. Problem images are observations; only " | |
| "reasoning images receive generation markers.\n", | |
| encoding="utf-8", | |
| ) | |
| return manifest | |
| def build_parser() -> argparse.ArgumentParser: | |
| parser = argparse.ArgumentParser(description=__doc__) | |
| parser.add_argument("--source-root", type=Path, default=Path("data/Zebra-CoT")) | |
| parser.add_argument("--output-dir", type=Path, default=Path("data/Zebra-CoT-samples")) | |
| parser.add_argument("--count", type=int, default=16) | |
| parser.add_argument("--image-size", type=int, default=256) | |
| parser.add_argument("--overwrite", action="store_true") | |
| return parser | |
| def main() -> None: | |
| args = build_parser().parse_args() | |
| manifest = prepare_samples( | |
| args.source_root, | |
| args.output_dir, | |
| count=args.count, | |
| image_size=args.image_size, | |
| overwrite=args.overwrite, | |
| ) | |
| print(json.dumps(manifest, indent=2, ensure_ascii=False)) | |
| if __name__ == "__main__": | |
| main() | |
Xet Storage Details
- Size:
- 7.28 kB
- Xet hash:
- ec0d617fda21f4e8169ce30ce6aeccfe9605218a6dbc80ecb0f83fd2970d3c59
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.