Buckets:
| """Create a deterministic per-config subset of raw Zebra-CoT Parquet files.""" | |
| from __future__ import annotations | |
| import argparse | |
| import hashlib | |
| import json | |
| import math | |
| import random | |
| import shutil | |
| import tempfile | |
| from bisect import bisect_left | |
| from pathlib import Path | |
| import pyarrow as pa | |
| import pyarrow.parquet as pq | |
| MANIFEST_NAME = "SPATIAL_MMCOT_SUBSET.json" | |
| def sample_count(total: int, fraction: float) -> int: | |
| """Round half up while retaining at least one row from a non-empty config.""" | |
| if total < 1: | |
| raise ValueError("cannot sample an empty config") | |
| if not 0.0 < fraction <= 1.0: | |
| raise ValueError("fraction must be in (0, 1]") | |
| return max(1, min(total, math.floor(total * fraction + 0.5))) | |
| def selected_indices(total: int, count: int, *, seed: int, config: str) -> list[int]: | |
| """Select stable global row indices without depending on filesystem ordering.""" | |
| digest = hashlib.sha256(f"{seed}:{config}".encode()).digest() | |
| rng = random.Random(int.from_bytes(digest[:8], "big")) | |
| return sorted(rng.sample(range(total), count)) | |
| def _schema_digest(schema: pa.Schema) -> str: | |
| return hashlib.sha256(schema.serialize().to_pybytes()).hexdigest() | |
| def _source_layout(files: list[Path]) -> tuple[list[int], list[int], pa.Schema]: | |
| rows = [] | |
| schemas = [] | |
| for path in files: | |
| parquet = pq.ParquetFile(path) | |
| rows.append(parquet.metadata.num_rows) | |
| schemas.append(parquet.schema_arrow) | |
| schema = schemas[0] | |
| if any(candidate != schema for candidate in schemas[1:]): | |
| raise ValueError(f"source schemas differ under {files[0].parent}") | |
| offsets = [0] | |
| for count in rows: | |
| offsets.append(offsets[-1] + count) | |
| return rows, offsets, schema | |
| def _selected_table(files: list[Path], indices: list[int], schema: pa.Schema) -> pa.Table: | |
| """Read only row groups containing selected global row indices.""" | |
| pieces: list[pa.Table] = [] | |
| global_file_start = 0 | |
| for path in files: | |
| parquet = pq.ParquetFile(path) | |
| file_end = global_file_start + parquet.metadata.num_rows | |
| first = bisect_left(indices, global_file_start) | |
| last = bisect_left(indices, file_end) | |
| local_indices = [index - global_file_start for index in indices[first:last]] | |
| selected_cursor = 0 | |
| row_group_start = 0 | |
| for row_group in range(parquet.num_row_groups): | |
| row_group_rows = parquet.metadata.row_group(row_group).num_rows | |
| row_group_end = row_group_start + row_group_rows | |
| begin = bisect_left(local_indices, row_group_start, selected_cursor) | |
| end = bisect_left(local_indices, row_group_end, begin) | |
| if begin != end: | |
| offsets = pa.array( | |
| [index - row_group_start for index in local_indices[begin:end]], | |
| type=pa.int64(), | |
| ) | |
| pieces.append(parquet.read_row_group(row_group).take(offsets)) | |
| selected_cursor = end | |
| row_group_start = row_group_end | |
| global_file_start = file_end | |
| if not pieces: | |
| raise RuntimeError("selection unexpectedly produced no rows") | |
| table = pa.concat_tables(pieces) | |
| if table.schema != schema: | |
| raise RuntimeError("selected table did not preserve the source schema") | |
| return table | |
| def create_subset( | |
| source_root: Path, | |
| output_root: Path, | |
| *, | |
| fraction: float = 0.01, | |
| seed: int = 20260825, | |
| ) -> dict: | |
| source_root = source_root.resolve() | |
| output_root = output_root.absolute() | |
| if output_root.exists() or output_root.is_symlink(): | |
| raise FileExistsError(f"output already exists: {output_root}") | |
| configs = sorted( | |
| path | |
| for path in source_root.iterdir() | |
| if path.is_dir() and path.name != ".git" and list(path.glob("*.parquet")) | |
| ) | |
| if not configs: | |
| raise FileNotFoundError(f"no Zebra-CoT config directories under {source_root}") | |
| output_root.parent.mkdir(parents=True, exist_ok=True) | |
| staging = Path(tempfile.mkdtemp(prefix=f".{output_root.name}-", dir=output_root.parent)) | |
| manifest = { | |
| "format": "zebra_cot_raw_parquet_subset_v1", | |
| "source_root": str(source_root), | |
| "fraction": fraction, | |
| "seed": seed, | |
| "rounding": "max(1, floor(source_rows * fraction + 0.5))", | |
| "configs": [], | |
| } | |
| try: | |
| for config_dir in configs: | |
| files = sorted(config_dir.glob("*.parquet")) | |
| source_rows, file_offsets, schema = _source_layout(files) | |
| total = sum(source_rows) | |
| count = sample_count(total, fraction) | |
| indices = selected_indices(total, count, seed=seed, config=config_dir.name) | |
| table = _selected_table(files, indices, schema) | |
| if table.num_rows != count: | |
| raise RuntimeError( | |
| f"{config_dir.name}: selected {table.num_rows} rows, expected {count}" | |
| ) | |
| destination_dir = staging / config_dir.name | |
| destination_dir.mkdir() | |
| destination = destination_dir / "train-00000-of-00001.parquet" | |
| pq.write_table(table, destination, compression="snappy", row_group_size=100) | |
| written = pq.ParquetFile(destination) | |
| if written.metadata.num_rows != count or written.schema_arrow != schema: | |
| raise RuntimeError(f"failed output verification for {config_dir.name}") | |
| selected_sources = [] | |
| for index in indices: | |
| file_index = bisect_left(file_offsets, index + 1) - 1 | |
| selected_sources.append( | |
| { | |
| "file": files[file_index].name, | |
| "row": index - file_offsets[file_index], | |
| "global_row": index, | |
| } | |
| ) | |
| manifest["configs"].append( | |
| { | |
| "name": config_dir.name, | |
| "source_rows": total, | |
| "sample_rows": count, | |
| "source_files": [ | |
| {"name": path.name, "rows": rows} | |
| for path, rows in zip(files, source_rows, strict=True) | |
| ], | |
| "schema_sha256": _schema_digest(schema), | |
| "output_file": str(destination.relative_to(staging)), | |
| "selected_rows": selected_sources, | |
| } | |
| ) | |
| print(f"SAMPLED {config_dir.name}: {count}/{total}", flush=True) | |
| manifest["source_rows"] = sum(item["source_rows"] for item in manifest["configs"]) | |
| manifest["sample_rows"] = sum(item["sample_rows"] for item in manifest["configs"]) | |
| (staging / MANIFEST_NAME).write_text( | |
| json.dumps(manifest, ensure_ascii=False, indent=2) + "\n" | |
| ) | |
| repair = source_root / "SPATIAL_MMCOT_REPAIR.json" | |
| if repair.exists(): | |
| shutil.copy2(repair, staging / repair.name) | |
| staging.rename(output_root) | |
| except BaseException: | |
| shutil.rmtree(staging, ignore_errors=True) | |
| raise | |
| return manifest | |
| def parse_args() -> argparse.Namespace: | |
| parser = argparse.ArgumentParser(description=__doc__) | |
| parser.add_argument("--source-root", type=Path, default=Path("data/Zebra-CoT")) | |
| parser.add_argument("--output-root", type=Path, default=Path("data/Zebra-CoT-small")) | |
| parser.add_argument("--fraction", type=float, default=0.01) | |
| parser.add_argument("--seed", type=int, default=20260825) | |
| return parser.parse_args() | |
| def main() -> None: | |
| args = parse_args() | |
| manifest = create_subset( | |
| args.source_root, | |
| args.output_root, | |
| fraction=args.fraction, | |
| seed=args.seed, | |
| ) | |
| print( | |
| f"DONE {manifest['sample_rows']}/{manifest['source_rows']} rows -> " | |
| f"{args.output_root.absolute()}" | |
| ) | |
| if __name__ == "__main__": | |
| main() | |
Xet Storage Details
- Size:
- 7.93 kB
- Xet hash:
- 8b0b3839abb839daff061613303c1ed185d3174a7d36c2b041b1586adb19a5b0
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.