Buckets:
| """Attach a deterministic balanced evaluation-subset index to a task manifest.""" | |
| from __future__ import annotations | |
| import argparse | |
| import hashlib | |
| import json | |
| import os | |
| import re | |
| import tempfile | |
| from collections import Counter, defaultdict | |
| from pathlib import Path | |
| import pyarrow.parquet as pq | |
| FORMAT = "zebra_cot_evaluation_subset_v1" | |
| FIELD = "evaluation_subset_index" | |
| def _digest(seed: int, namespace: str, sample_id: str) -> str: | |
| return hashlib.sha256(f"{seed}:{namespace}:{sample_id}".encode()).hexdigest() | |
| def _choice(value: str) -> str: | |
| match = re.fullmatch(r"\s*([A-Da-d])(?:[.)])?\s*", value) | |
| if match is None: | |
| raise ValueError(f"expected an A-D answer, got {value!r}") | |
| return match.group(1).upper() | |
| def _targets(task_root: Path, rows: list[dict]) -> dict[str, str]: | |
| by_shard: dict[str, list[dict]] = defaultdict(list) | |
| for row in rows: | |
| by_shard[row["source"]["shard"]].append(row) | |
| targets = {} | |
| for relative, shard_rows in by_shard.items(): | |
| table = pq.read_table(task_root / relative, columns=["Final Answer"]) | |
| column = table.column("Final Answer") | |
| for row in shard_rows: | |
| targets[row["sample_id"]] = column[row["source"]["row"]].as_py() | |
| return targets | |
| def attach_subset_index( | |
| cache_root: Path, | |
| *, | |
| per_choice: int = 25, | |
| seed: int = 20260827, | |
| ) -> dict: | |
| """Select balanced test rows and atomically rewrite the cache manifest.""" | |
| cache_root = cache_root.resolve() | |
| task_root = cache_root.parent | |
| manifest = cache_root / "manifest.jsonl" | |
| metadata_path = task_root / "SPATIAL_MMCOT_EVALUATION_SUBSET.json" | |
| if metadata_path.exists(): | |
| raise FileExistsError(f"evaluation subset already exists: {metadata_path}") | |
| rows = [json.loads(line) for line in manifest.read_text().splitlines() if line] | |
| if any(FIELD in row for row in rows): | |
| raise ValueError(f"manifest already contains {FIELD}") | |
| test_rows = [row for row in rows if row["split"] == "test"] | |
| targets = _targets(task_root, test_rows) | |
| candidates: dict[str, list[dict]] = defaultdict(list) | |
| for row in test_rows: | |
| candidates[_choice(targets[row["sample_id"]])].append(row) | |
| if set(candidates) != set("ABCD"): | |
| raise ValueError(f"expected A-D buckets, got {sorted(candidates)}") | |
| selected = [] | |
| for choice in "ABCD": | |
| bucket = sorted( | |
| candidates[choice], | |
| key=lambda row: _digest(seed, f"select-{choice}", row["sample_id"]), | |
| ) | |
| if len(bucket) < per_choice: | |
| raise ValueError(f"choice {choice} has only {len(bucket)} candidates") | |
| selected.extend(bucket[:per_choice]) | |
| selected.sort(key=lambda row: _digest(seed, "order", row["sample_id"])) | |
| indices = {row["sample_id"]: index for index, row in enumerate(selected)} | |
| for row in rows: | |
| row[FIELD] = indices.get(row["sample_id"]) | |
| handle, temporary_name = tempfile.mkstemp( | |
| prefix=".manifest-eval-", suffix=".jsonl", dir=cache_root | |
| ) | |
| try: | |
| with os.fdopen(handle, "w", encoding="utf-8") as output: | |
| for row in rows: | |
| output.write(json.dumps(row, ensure_ascii=False) + "\n") | |
| output.flush() | |
| os.fsync(output.fileno()) | |
| os.chmod(temporary_name, manifest.stat().st_mode) | |
| Path(temporary_name).replace(manifest) | |
| except BaseException: | |
| Path(temporary_name).unlink(missing_ok=True) | |
| raise | |
| selected_metadata = [ | |
| { | |
| "index": indices[row["sample_id"]], | |
| "sample_id": row["sample_id"], | |
| "choice": _choice(targets[row["sample_id"]]), | |
| "source": row["source"], | |
| "sequence_length": row["sequence_length"], | |
| "generated_images": row["generated_images"], | |
| "visual_positions": row["visual_positions"], | |
| } | |
| for row in selected | |
| ] | |
| selected_metadata.sort(key=lambda row: row["index"]) | |
| summary = { | |
| "format": FORMAT, | |
| "manifest": str(manifest), | |
| "field": FIELD, | |
| "split": "test", | |
| "seed": seed, | |
| "selection": "lowest seeded SHA256 independently within each answer choice", | |
| "samples": len(selected), | |
| "per_choice": dict(Counter(row["choice"] for row in selected_metadata)), | |
| "rows": selected_metadata, | |
| } | |
| metadata_path.write_text( | |
| json.dumps(summary, ensure_ascii=False, indent=2) + "\n", encoding="utf-8" | |
| ) | |
| return summary | |
| def main() -> None: | |
| parser = argparse.ArgumentParser(description=__doc__) | |
| parser.add_argument("--cache-root", type=Path, required=True) | |
| parser.add_argument("--per-choice", type=int, default=25) | |
| parser.add_argument("--seed", type=int, default=20260827) | |
| args = parser.parse_args() | |
| if args.per_choice <= 0: | |
| raise ValueError("per-choice must be positive") | |
| print( | |
| json.dumps( | |
| attach_subset_index( | |
| args.cache_root, per_choice=args.per_choice, seed=args.seed | |
| ), | |
| ensure_ascii=False, | |
| indent=2, | |
| ) | |
| ) | |
| if __name__ == "__main__": | |
| main() | |
Xet Storage Details
- Size:
- 5.14 kB
- Xet hash:
- 1c897eb3b3810b43afe94afe9c9e7c152ba8badaae52cd6e4e1914dc3425d9b7
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.