Buckets:
| """Fixed, image-grouped 95/5 development view; the all-train cache stays intact.""" | |
| from __future__ import annotations | |
| import argparse | |
| from collections import Counter, defaultdict | |
| from concurrent.futures import ProcessPoolExecutor | |
| import hashlib | |
| import io | |
| import json | |
| import math | |
| import multiprocessing | |
| from pathlib import Path | |
| import re | |
| import tempfile | |
| import numpy as np | |
| from PIL import Image, ImageOps | |
| TASKS = ("Chart_Refocus", "Jigsaw_Assembly", "Spatial_Navigation", "Visual_Search") | |
| EXPECTED_COUNTS = dict(zip(TASKS, (6000, 6000, 6000, 6990))) | |
| FORMAT = "thinkmorph_grouped_id_validation_v1" | |
| def file_hash(path): | |
| digest = hashlib.sha256() | |
| with Path(path).open("rb") as stream: | |
| for block in iter(lambda: stream.read(8 * 1024 * 1024), b""): | |
| digest.update(block) | |
| return digest.hexdigest() | |
| def image_fingerprint(payload): | |
| with Image.open(io.BytesIO(payload)) as source: | |
| image = ImageOps.exif_transpose(source).convert("RGB") | |
| digest = hashlib.sha256(str(image.size).encode() + b"\0" + image.tobytes()).hexdigest() | |
| horizontal = np.asarray(image.convert("L").resize((9, 8), Image.Resampling.LANCZOS)) | |
| vertical = np.asarray(image.convert("L").resize((8, 9), Image.Resampling.LANCZOS)) | |
| bits = np.concatenate(((horizontal[:, 1:] > horizontal[:, :-1]).ravel(), | |
| (vertical[1:, :] > vertical[:-1, :]).ravel())) | |
| dhash = int.from_bytes(np.packbits(bits).tobytes(), "big") | |
| thumbnail = image.resize((32, 32), Image.Resampling.LANCZOS).tobytes() | |
| return {"sha256": digest, "size": image.size, "dhash": dhash, "thumbnail": thumbnail} | |
| def near_duplicate(first, second): | |
| if (first["dhash"] ^ second["dhash"]).bit_count() > 2: | |
| return False | |
| a, b = first["size"], second["size"] | |
| if abs((a[0] / a[1]) / (b[0] / b[1]) - 1) > .001: | |
| return False | |
| diff = np.abs(np.frombuffer(first["thumbnail"], np.uint8).astype(np.int16) | |
| - np.frombuffer(second["thumbnail"], np.uint8).astype(np.int16)) | |
| return float(diff.mean()) <= 1.5 and float(np.mean(diff > 8)) <= .01 | |
| def source_identity(task, path): | |
| """Known original-image families; never group questions by answer/template.""" | |
| stem = Path(path or "").stem | |
| if not stem: | |
| return None | |
| if task == "Jigsaw_Assembly": | |
| stem = re.sub(r"_(border|answer_\d+x\d+)$", "", stem) | |
| elif task == "Spatial_Navigation": | |
| stem = re.sub(r"(_cot)?(_visualization)?$", "", stem) | |
| return f"{task}/{stem}" | |
| def subtype(task, row): | |
| path = row["problem_image_0"].get("path") or "" | |
| target = row["reasoning_image_0"].get("path") or "" | |
| if task == "Jigsaw_Assembly": | |
| match = re.search(r"answer_(\d+x\d+)", target) | |
| return match[1] if match else "unknown_grid" | |
| if task == "Spatial_Navigation": | |
| match = re.search(r"frozenlake_s(\d+)_", path) | |
| return "grid_" + match[1] if match else "unknown_grid" | |
| if task == "Chart_Refocus": | |
| return "two_col" if path.startswith("two_col_") else "other_chart" | |
| return "gqa" if "gqa_" in target else "other_search" | |
| def fingerprint_row(item): | |
| row, raw = item | |
| task = row["config"] | |
| images = [] | |
| for key in ("problem_image_0", "reasoning_image_0"): | |
| cell = raw[key] | |
| images.append(dict(image_fingerprint(cell["bytes"]), role=key, path=cell.get("path"))) | |
| # The generated picture is also compared: different permutations of one | |
| # jigsaw can have different inputs but the same reconstructed original. | |
| source_keys = [source_identity(task, images[0]["path"])] | |
| if task in ("Jigsaw_Assembly", "Spatial_Navigation"): | |
| source_keys.append(source_identity(task, images[1]["path"])) | |
| return {"sample_id": row["sample_id"], "task": task, "subtype": subtype(task, raw), | |
| "source_keys": sorted({key for key in source_keys if key}), "images": images} | |
| def group_rows(records): | |
| parent = list(range(len(records))) | |
| def root(index): | |
| while parent[index] != index: | |
| parent[index] = parent[parent[index]] | |
| index = parent[index] | |
| return index | |
| def join(first, second): | |
| a, b = root(first), root(second) | |
| if a != b: | |
| parent[b] = a | |
| identities, exact, bands = {}, {}, defaultdict(list) | |
| edges = [] | |
| for index, record in enumerate(records): | |
| for key in record["source_keys"]: | |
| if key in identities: | |
| join(index, identities[key]) | |
| identities[key] = index | |
| for fingerprint in record["images"]: | |
| key = fingerprint["sha256"] | |
| if key in exact: | |
| join(index, exact[key]) | |
| exact[key] = index | |
| # Natural-image re-encodings/resizes only. Sparse charts/maze | |
| # templates are grouped by source identity and exact decoded pixels, | |
| # not a perceptual threshold that can merge distinct diagrams. | |
| if record["task"] not in ("Jigsaw_Assembly", "Visual_Search"): | |
| continue | |
| candidates = {} | |
| for band in range(4): | |
| band_key = (band, (fingerprint["dhash"] >> (band * 32)) & (2**32 - 1)) | |
| for other, prior in bands[band_key]: | |
| candidates[(other, prior["sha256"])] = prior | |
| for (other, _), prior in candidates.items(): | |
| if root(index) != root(other) and near_duplicate(fingerprint, prior): | |
| join(index, other) | |
| edges.append((record["sample_id"], records[other]["sample_id"])) | |
| for band in range(4): | |
| bands[(band, (fingerprint["dhash"] >> (band * 32)) & (2**32 - 1))].append((index, fingerprint)) | |
| members = defaultdict(list) | |
| for index, record in enumerate(records): | |
| members[root(index)].append(record["sample_id"]) | |
| assignments = {} | |
| for values in members.values(): | |
| group = hashlib.sha256("\n".join(sorted(values)).encode()).hexdigest() | |
| assignments.update((sample, group) for sample in values) | |
| return assignments, edges | |
| def select_groups(groups, target, seed): | |
| """Deterministic whole-group subset closest to the requested row count.""" | |
| if not groups: | |
| return set() | |
| order = sorted(groups, key=lambda group: hashlib.sha256(f"{seed}\0{group}".encode()).digest()) | |
| limit = target + max(len(groups[group]) for group in groups) | |
| reachable = {0: None} | |
| for group in order: | |
| weight = len(groups[group]) | |
| for size in sorted(reachable, reverse=True): | |
| if size + weight <= limit and size + weight not in reachable: | |
| reachable[size + weight] = (size, group) | |
| if target in reachable: | |
| break | |
| size = min(reachable, key=lambda size: (abs(size - target), size > target)) | |
| selected = set() | |
| while size: | |
| size, group = reachable[size] | |
| selected.add(group) | |
| return selected | |
| def split_records(records, *, fraction=.05, seed=17): | |
| if not 0 < fraction < 1: | |
| raise ValueError("validation fraction must lie between zero and one") | |
| groups, edges = group_rows(records) | |
| grouped = defaultdict(list) | |
| for record in records: | |
| grouped[groups[record["sample_id"]]].append(record) | |
| strata = defaultdict(dict) | |
| for group, members in grouped.items(): | |
| # Cross-task / cross-difficulty groups stay intact in a joint stratum. | |
| signature = tuple(sorted({(row["task"], row["subtype"]) for row in members})) | |
| strata[signature][group] = members | |
| held = set() | |
| for signature, candidates in sorted(strata.items()): | |
| target = math.floor(sum(map(len, candidates.values())) * fraction + .5) | |
| held.update(select_groups(candidates, target, seed)) | |
| assignments = {sample: (group, "validation" if group in held else "train") for sample, group in groups.items()} | |
| counts, subtype_counts = defaultdict(Counter), defaultdict(Counter) | |
| for record in records: | |
| split = assignments[record["sample_id"]][1] | |
| counts[record["task"]][split] += 1 | |
| subtype_counts[f"{record['task']}/{record['subtype']}"][split] += 1 | |
| if any(not counts[task][split] for task in counts for split in ("train", "validation")): | |
| raise ValueError("grouping leaves a task without train or validation data") | |
| return assignments, {"task_counts": {task: dict(c) for task, c in counts.items()}, | |
| "subtype_counts": {key: dict(c) for key, c in sorted(subtype_counts.items())}, | |
| "groups": len(grouped), "largest_group": max(map(len, grouped.values())), | |
| "multirow_groups": sum(len(v) > 1 for v in grouped.values()), | |
| "near_duplicate_edges": edges} | |
| def audit_split(records, assignments): | |
| split_sets = defaultdict(set) | |
| for row in records: | |
| group, split = assignments[row["sample_id"]] | |
| for key in ["group:" + group, *row["source_keys"], | |
| *("pixels:" + f["sha256"] for f in row["images"])]: | |
| split_sets[key].add(split) | |
| if any(len(values) > 1 for values in split_sets.values()): | |
| raise AssertionError("source/image/group leakage across train and validation") | |
| def build_view(source_root, source_cache, output, *, seed=17, workers=2): | |
| import pyarrow.parquet as pq | |
| source_root, source_cache, output = map(lambda p: Path(p).resolve(), (source_root, source_cache, output)) | |
| if output.exists(): | |
| raise FileExistsError(f"fixed split already exists; reuse it: {output}") | |
| if not (source_cache / "READY").is_file(): | |
| raise ValueError("source cache is incomplete") | |
| manifest = source_cache / "manifest.jsonl" | |
| manifest_hash = file_hash(manifest) | |
| rows = [json.loads(line) for line in manifest.read_text().splitlines() if line] | |
| if dict(Counter(row["config"] for row in rows)) != EXPECTED_COUNTS: | |
| raise ValueError("unexpected ThinkMorph task counts") | |
| if any(row["split"] != "train" for row in rows): | |
| raise ValueError("development view must originate from the all-train cache") | |
| by_source = {(row["source"]["shard"], int(row["source"]["row"])): row for row in rows} | |
| if len(by_source) != len(rows) or len({r["sample_id"] for r in rows}) != len(rows): | |
| raise ValueError("duplicate cache coordinates/sample IDs") | |
| fingerprints = [] | |
| raw_hashes = {} | |
| with ProcessPoolExecutor(max_workers=workers, mp_context=multiprocessing.get_context("spawn")) as pool: | |
| for relative in sorted({key[0] for key in by_source}): | |
| path = source_root / relative | |
| before = (path.stat().st_size, path.stat().st_mtime_ns) | |
| raw_hashes[relative] = file_hash(path) | |
| offset = 0 | |
| for batch in pq.ParquetFile(path).iter_batches(batch_size=32, columns=["pid", "problem_image_0", "reasoning_image_0"]): | |
| work = [] | |
| for raw in batch.to_pylist(): | |
| row = by_source[(relative, offset)] | |
| if str(raw["pid"]) != str(row["source"]["pid"]): | |
| raise ValueError("raw/cache pid mismatch") | |
| work.append((row, raw)) | |
| offset += 1 | |
| fingerprints.extend(pool.map(fingerprint_row, work, chunksize=4)) | |
| if before != (path.stat().st_size, path.stat().st_mtime_ns): | |
| raise RuntimeError("raw source changed during split creation") | |
| print(f"SPLIT_FINGERPRINTED {relative} total={len(fingerprints)}", flush=True) | |
| if len(fingerprints) != len(rows) or file_hash(manifest) != manifest_hash: | |
| raise RuntimeError("source cache changed or raw coverage is incomplete") | |
| fingerprints.sort(key=lambda row: row["sample_id"]) | |
| assignments, report = split_records(fingerprints, seed=seed) | |
| audit_split(fingerprints, assignments) | |
| for left, right in report["near_duplicate_edges"]: | |
| assert assignments[left][1] == assignments[right][1] | |
| report.update(format=FORMAT, seed=seed, validation_fraction=.05, source_cache=str(source_cache), | |
| source_manifest_sha256=manifest_hash, raw_sha256=raw_hashes, | |
| rows=len(rows), counts=dict(Counter(split for _, split in assignments.values())), | |
| group_leakage=0, input_or_target_pixel_leakage=0, source_identity_leakage=0, | |
| method="whole connected image/source groups; task/subtype strata; seeded nearest-row-count subset", | |
| near_duplicate_policy="natural images only: 128-bit dHash distance<=2; aspect ratio within 0.1%; 32px RGB MAE<=1.5/255, fraction absdiff>8/255<=1%", | |
| limitation="not an exhaustive detector of arbitrary crops or semantically identical rerenders", | |
| purpose="development validation only; final training uses the untouched all-train cache") | |
| output.parent.mkdir(parents=True, exist_ok=True) | |
| with tempfile.TemporaryDirectory(prefix=".idval5-", dir=output.parent) as temporary: | |
| staging = Path(temporary) / "view" | |
| staging.mkdir() | |
| with (staging / "manifest.jsonl").open("w") as stream: | |
| for original in rows: | |
| group, split = assignments[original["sample_id"]] | |
| row = dict(original, source_group_id=original["group_id"], group_id=group, | |
| split=split, split_score=int(group[:16], 16) / 2**64, | |
| cache_shard=str((source_cache / original["cache_shard"]).resolve())) | |
| stream.write(json.dumps(row, ensure_ascii=False) + "\n") | |
| with (staging / "image_groups.jsonl").open("w") as stream: | |
| for record in fingerprints: | |
| value = {**record, "group_id": assignments[record["sample_id"]][0], "split": assignments[record["sample_id"]][1]} | |
| value["images"] = [{k: v for k, v in image.items() if k != "thumbnail"} for image in record["images"]] | |
| stream.write(json.dumps(value, ensure_ascii=False) + "\n") | |
| config = json.loads((source_cache / "config.json").read_text()) | |
| config["split_policy"] = FORMAT | |
| config["development_view"] = {"source_cache": str(source_cache), "seed": seed, "validation_fraction": .05} | |
| report["manifest_sha256"] = file_hash(staging / "manifest.jsonl") | |
| report["image_groups_sha256"] = file_hash(staging / "image_groups.jsonl") | |
| for filename, value in (("config.json", config), ("SPLIT.json", report)): | |
| (staging / filename).write_text(json.dumps(value, ensure_ascii=False, indent=2) + "\n") | |
| (staging / "READY").write_text("complete\n") | |
| if output.exists(): | |
| raise FileExistsError(output) | |
| staging.rename(output) | |
| return report | |
| def main(): | |
| parser = argparse.ArgumentParser(description=__doc__) | |
| parser.add_argument("--source-root", type=Path, default=Path("data/ThinkMorph")) | |
| parser.add_argument("--source-cache", type=Path, default=Path("data/ThinkMorph/ThinkMorph")) | |
| parser.add_argument("--output", type=Path, default=Path("data/ThinkMorph/ThinkMorph-IDVal5")) | |
| parser.add_argument("--seed", type=int, default=17) | |
| parser.add_argument("--workers", type=int, default=2) | |
| args = parser.parse_args() | |
| report = build_view(args.source_root, args.source_cache, args.output, seed=args.seed, workers=args.workers) | |
| print(json.dumps({k: v for k, v in report.items() if k != "near_duplicate_edges"}), flush=True) | |
| if __name__ == "__main__": | |
| main() | |
Xet Storage Details
- Size:
- 15.5 kB
- Xet hash:
- df36b71037028e967d40c6b2c71af027d623a28220ea7199084cf34f98f257bf
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.