Buckets:
| """Build zero-copy, group-safe task views of the packed Zebra-CoT cache.""" | |
| from __future__ import annotations | |
| import argparse | |
| import hashlib | |
| import json | |
| import tempfile | |
| from collections import Counter | |
| from pathlib import Path | |
| FORMAT = "zebra_cot_task_split_v1" | |
| def _sha256(path: Path) -> str: | |
| digest = hashlib.sha256() | |
| with path.open("rb") as handle: | |
| for chunk in iter(lambda: handle.read(1024 * 1024), b""): | |
| digest.update(chunk) | |
| return digest.hexdigest() | |
| def build_task_view( | |
| source_cache: Path, | |
| task_root: Path, | |
| config_name: str, | |
| *, | |
| validation_fraction: float = 0.05, | |
| test_fraction: float = 0.05, | |
| ) -> dict: | |
| """Materialize only metadata; tensor shards remain in the shared cache.""" | |
| if validation_fraction <= 0 or test_fraction <= 0: | |
| raise ValueError("validation/test fractions must be positive") | |
| if validation_fraction + test_fraction >= 1: | |
| raise ValueError("validation + test fractions must be below one") | |
| source_cache = source_cache.resolve() | |
| task_root = task_root.resolve() | |
| source_manifest = source_cache / "manifest.jsonl" | |
| raw_metadata = task_root / "SPATIAL_MMCOT_TASK_SUBSET.json" | |
| if not source_manifest.is_file() or not (source_cache / "READY").is_file(): | |
| raise FileNotFoundError(f"incomplete source cache: {source_cache}") | |
| if not raw_metadata.is_file(): | |
| raise FileNotFoundError( | |
| f"missing task materialization metadata: {raw_metadata}" | |
| ) | |
| raw = json.loads(raw_metadata.read_text(encoding="utf-8")) | |
| if raw["source_config"] != config_name: | |
| raise ValueError("task root and requested config disagree") | |
| rows = [] | |
| with source_manifest.open(encoding="utf-8") as handle: | |
| for line in handle: | |
| if not line.strip(): | |
| continue | |
| row = json.loads(line) | |
| if row["config"] == config_name: | |
| rows.append(row) | |
| if len(rows) != raw["rows"] or any(row["config"] != config_name for row in rows): | |
| raise ValueError( | |
| f"packed cache has {len(rows)} rows for {config_name}, expected {raw['rows']}" | |
| ) | |
| groups = sorted({row["group_id"] for row in rows}) | |
| test_groups = round(len(groups) * test_fraction) | |
| validation_groups = round(len(groups) * validation_fraction) | |
| if test_groups + validation_groups >= len(groups): | |
| raise ValueError("not enough groups for requested splits") | |
| assignments = { | |
| group: ( | |
| "test" | |
| if index < test_groups | |
| else "validation" | |
| if index < test_groups + validation_groups | |
| else "train" | |
| ) | |
| for index, group in enumerate(groups) | |
| } | |
| split_rows = [] | |
| for row in rows: | |
| row["split"] = assignments[row["group_id"]] | |
| row["cache_shard"] = str((source_cache / row["cache_shard"]).resolve()) | |
| split_rows.append(row) | |
| group_splits: dict[str, set[str]] = {} | |
| for row in split_rows: | |
| group_splits.setdefault(row["group_id"], set()).add(row["split"]) | |
| if any(len(splits) != 1 for splits in group_splits.values()): | |
| raise RuntimeError("group leakage in task split") | |
| row_counts = Counter(row["split"] for row in split_rows) | |
| group_counts = Counter(assignments.values()) | |
| summary = { | |
| "format": FORMAT, | |
| "source_cache": str(source_cache), | |
| "source_manifest_sha256": _sha256(source_manifest), | |
| "config": config_name, | |
| "method": "lexicographically sorted SHA256 group_id; whole-group assignment", | |
| "requested_fractions": { | |
| "train": 1.0 - validation_fraction - test_fraction, | |
| "validation": validation_fraction, | |
| "test": test_fraction, | |
| }, | |
| "rows": dict(row_counts), | |
| "groups": dict(group_counts), | |
| "total_rows": len(split_rows), | |
| "total_groups": len(groups), | |
| "group_leakage": 0, | |
| } | |
| output = task_root / "UniAR-960-cache" | |
| split_manifest = task_root / "SPATIAL_MMCOT_SPLIT.json" | |
| if output.exists() or split_manifest.exists(): | |
| raise FileExistsError(f"task split output already exists under {task_root}") | |
| staging = Path(tempfile.mkdtemp(prefix=".UniAR-960-cache-", dir=task_root)) | |
| try: | |
| with (staging / "manifest.jsonl").open("w", encoding="utf-8") as handle: | |
| for row in split_rows: | |
| handle.write(json.dumps(row, ensure_ascii=False) + "\n") | |
| source_config = json.loads((source_cache / "config.json").read_text()) | |
| source_config["task_view"] = summary | |
| (staging / "config.json").write_text( | |
| json.dumps(source_config, ensure_ascii=False, indent=2) + "\n", | |
| encoding="utf-8", | |
| ) | |
| (staging / "READY").write_text("complete\n", encoding="utf-8") | |
| staging.rename(output) | |
| split_manifest.write_text( | |
| json.dumps(summary, ensure_ascii=False, indent=2) + "\n", | |
| encoding="utf-8", | |
| ) | |
| except BaseException: | |
| if staging.exists(): | |
| for path in staging.iterdir(): | |
| path.unlink() | |
| staging.rmdir() | |
| raise | |
| return summary | |
| def main() -> None: | |
| parser = argparse.ArgumentParser(description=__doc__) | |
| parser.add_argument("--source-cache", type=Path, required=True) | |
| parser.add_argument("--task-root", type=Path, required=True) | |
| parser.add_argument("--config", required=True) | |
| parser.add_argument("--validation-fraction", type=float, default=0.05) | |
| parser.add_argument("--test-fraction", type=float, default=0.05) | |
| args = parser.parse_args() | |
| summary = build_task_view( | |
| args.source_cache, | |
| args.task_root, | |
| args.config, | |
| validation_fraction=args.validation_fraction, | |
| test_fraction=args.test_fraction, | |
| ) | |
| print(json.dumps(summary, ensure_ascii=False, indent=2)) | |
| if __name__ == "__main__": | |
| main() | |
Xet Storage Details
- Size:
- 5.93 kB
- Xet hash:
- a683696601dfde1d6116005ba55ace0cf59c2d55f3cd0bfdf14c1504c249f754
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.