Buckets:

EmergentTwilight
/

MMCoT

Files

xet

EmergentTwilight/MMCoT / experiments /src /thinkmorph_sft /prepare_cache.py

EmergentTwilight

21 days ago

download

raw

14.8 kB

"""Build a train-only canonical UniAR cache from the 24,990 ThinkMorph rows."""

from __future__ import annotations

import argparse
import json
import os
import time
from collections import Counter, defaultdict
from datetime import timedelta
from pathlib import Path

import pyarrow.parquet as pq
import torch
import torch.distributed as dist
from sft.runtime import file_sha256, source_fingerprint, write_json
from transformers import AutoProcessor
from uniar import UniARConfig, configure_processor
from uniar.vision import load_bsq_image_tokenizer
from zebra_sft.full_data import (
    CACHE_SCHEMA,
    IMAGE_PREPROCESSING_PROFILES,
    LEGACY_IMAGE_PREPROCESSING,
    prepare_row,
)
from zebra_sft.prepare_full import CacheWriter
from zebra_sft.protocol import SERIALIZATION_VERSION, SYSTEM_PROMPT


TASK_COUNTS = {
    "Chart_Refocus": 6000,
    "Jigsaw_Assembly": 6000,
    "Spatial_Navigation": 6000,
    "Visual_Search": 6990,
}


def source_shards(root: Path) -> list[Path]:
    paths = [
        path
        for task in sorted(TASK_COUNTS)
        for path in sorted((root / task / "data").glob("train-*.parquet"))
    ]
    if not paths:
        raise FileNotFoundError(f"no ThinkMorph training parquet under {root}")
    return paths


def canonical_source_row(row: dict) -> dict:
    """Map one public ThinkMorph row to the canonical typed trajectory inputs."""

    required_text = (
        "pid",
        "question",
        "answer",
        "resoning_thought_0",
        "resoning_thought_1",
        "task",
    )
    for key in required_text:
        if not isinstance(row.get(key), str) or not row[key].strip():
            raise ValueError(f"ThinkMorph row has empty {key}")
    for key in ("problem_image_0", "reasoning_image_0"):
        if not isinstance(row.get(key), dict) or not row[key].get("bytes"):
            raise ValueError(f"ThinkMorph row has no bytes for {key}")
    return {
        "Question": (
            f"{row['question'].strip()}\n"
            "<image_start>[problem_image_0]<image_end>"
        ),
        "Text Reasoning Trace": (
            f"{row['resoning_thought_0'].strip()}\n"
            "<image_start>[reasoning_image_0]<image_end>\n"
            f"{row['resoning_thought_1'].strip()}"
        ),
        "Final Answer": row["answer"].strip(),
        "problem_image_0": row["problem_image_0"],
        "reasoning_image_0": row["reasoning_image_0"],
    }


def _mode(output_root: Path, auto_resume: bool) -> str:
    if (output_root / "READY").is_file():
        return "complete"
    if not output_root.exists() or not any(output_root.iterdir()):
        return "fresh"
    if auto_resume:
        return "resume"
    raise FileExistsError(f"cache output already exists: {output_root}")


def _read_cached_coordinates(manifest_root: Path) -> set[tuple[str, int]]:
    coordinates = set()
    for path in sorted(manifest_root.glob("rank-*.jsonl")):
        data = path.read_bytes()
        boundary = data.rfind(b"\n") + 1
        if boundary != len(data):
            with path.open("r+b") as handle:
                handle.truncate(boundary)
            data = data[:boundary]
        for line in data.decode("utf-8").splitlines():
            if not line:
                continue
            row = json.loads(line)
            coordinate = (row["source"]["shard"], int(row["source"]["row"]))
            if coordinate in coordinates:
                raise RuntimeError(f"duplicate cached coordinate: {coordinate}")
            coordinates.add(coordinate)
    return coordinates


def _assign_shards(
    shards: list[Path], root: Path, cached: set[tuple[str, int]], world_size: int
) -> list[list[int]]:
    cached_counts = Counter(shard for shard, _ in cached)
    pending = []
    for index, path in enumerate(shards):
        relative = str(path.relative_to(root))
        remaining = pq.ParquetFile(path).metadata.num_rows - cached_counts[relative]
        if remaining > 0:
            pending.append((remaining, index))
    assignments = [[] for _ in range(world_size)]
    loads = [0] * world_size
    for remaining, index in sorted(pending, reverse=True):
        rank = min(range(world_size), key=lambda value: (loads[value], value))
        assignments[rank].append(index)
        loads[rank] += remaining
    return assignments


def _aggregate(output_root: Path, world_size: int, config: dict) -> dict:
    rows = []
    rank_stats = []
    for rank in range(world_size):
        manifest = output_root / "manifests" / f"rank-{rank:02d}.jsonl"
        rows.extend(
            json.loads(line)
            for line in manifest.read_text(encoding="utf-8").splitlines()
            if line
        )
        rank_stats.append(
            json.loads(
                (output_root / "stats" / f"rank-{rank:02d}.json").read_text(
                    encoding="utf-8"
                )
            )
        )
    rows.sort(key=lambda row: (row["config"], row["source"]["shard"], row["source"]["row"]))
    if len(rows) != sum(TASK_COUNTS.values()):
        raise RuntimeError(f"expected 24990 ThinkMorph rows, got {len(rows)}")
    if len({row["sample_id"] for row in rows}) != len(rows):
        raise RuntimeError("duplicate stable sample IDs in ThinkMorph cache")
    task_counts = Counter(row["config"] for row in rows)
    if dict(task_counts) != TASK_COUNTS:
        raise RuntimeError(f"ThinkMorph task counts differ: {dict(task_counts)}")
    for row in rows:
        row["split"] = "train"

    temporary = output_root / ".manifest.jsonl.tmp"
    with temporary.open("w", encoding="utf-8") as handle:
        for row in rows:
            handle.write(json.dumps(row, ensure_ascii=False) + "\n")
    temporary.replace(output_root / "manifest.jsonl")
    stats = {
        "schema": CACHE_SCHEMA,
        "config": config,
        "samples": len(rows),
        "splits": {"train": len(rows)},
        "config_splits": {
            task: {"train": count} for task, count in sorted(task_counts.items())
        },
        "unique_groups": len({row["group_id"] for row in rows}),
        "group_leakage": 0,
        "selected_cap": dict(Counter(str(row["selected_cap"]) for row in rows)),
        "sequence_length": {
            "max": max(row["sequence_length"] for row in rows),
            "mean": sum(row["sequence_length"] for row in rows) / len(rows),
        },
        "visual_positions": {
            "max": max(row["visual_positions"] for row in rows),
            "mean": sum(row["visual_positions"] for row in rows) / len(rows),
        },
        "rank_stats": rank_stats,
    }
    write_json(output_root / "stats.json", stats)
    ready = output_root / ".READY.tmp"
    ready.write_text("complete\n", encoding="utf-8")
    ready.replace(output_root / "READY")
    return stats


def prepare(args: argparse.Namespace) -> None:
    mode = _mode(args.output_root, args.auto_resume)
    if mode == "complete":
        print(f"CACHE_ALREADY_COMPLETE {args.output_root}", flush=True)
        return
    if not torch.cuda.is_available():
        raise RuntimeError("ThinkMorph cache preparation requires CUDA")
    local_rank = int(os.environ.get("LOCAL_RANK", 0))
    torch.cuda.set_device(local_rank)
    device = torch.device("cuda", local_rank)
    dist.init_process_group("nccl", device_id=device, timeout=timedelta(hours=6))
    rank, world_size = dist.get_rank(), dist.get_world_size()

    shards = source_shards(args.source_root)
    config = {
        "schema": "thinkmorph_uniar_cache_prepare_v1",
        "source_root": str(args.source_root.resolve()),
        "model_path": str(args.model_path.resolve()),
        "source_fingerprint": source_fingerprint("shared", "sft", "uniar"),
        "producer_sha256": file_sha256(Path(__file__)),
        "full_data_sha256": file_sha256(
            Path(__file__).parents[1] / "zebra_sft" / "full_data.py"
        ),
        "trajectory_sha256": file_sha256(
            Path(__file__).parents[1] / "zebra_sft" / "trajectory.py"
        ),
        "serialization_version": SERIALIZATION_VERSION,
        "system_prompt": SYSTEM_PROMPT,
        "split_policy": "all_public_rows_are_train",
        "task_counts": TASK_COUNTS,
        "max_long_edge": args.max_long_edge,
        "image_preprocessing_profile": args.image_preprocessing_profile,
        "target_canvas_size": args.target_canvas_size,
        "target_background": args.target_background,
        "max_sequence_length": args.max_sequence_length,
        "shard_size": args.shard_size,
        "world_size": world_size,
        "source_shards": len(shards),
    }
    if rank == 0:
        if mode == "fresh":
            args.output_root.mkdir(parents=True, exist_ok=True)
            (args.output_root / "manifests").mkdir()
            (args.output_root / "stats").mkdir()
            write_json(args.output_root / "config.json", config)
        else:
            previous = json.loads(
                (args.output_root / "config.json").read_text(encoding="utf-8")
            )
            if previous != config:
                raise ValueError("ThinkMorph cache resume configuration mismatch")
    dist.barrier()

    processor = configure_processor(
        AutoProcessor.from_pretrained(args.model_path, use_fast=True)
    )
    model_config = UniARConfig.from_pretrained(args.model_path)
    model = load_bsq_image_tokenizer(
        str(args.model_path),
        subfolder="bsq_encoder",
        attn_implementation=args.attn,
    ).to(device).eval()

    cached = (
        _read_cached_coordinates(args.output_root / "manifests")
        if mode == "resume"
        else set()
    )
    assignments = _assign_shards(shards, args.source_root, cached, world_size)
    owned = assignments[rank]
    writer = CacheWriter(args.output_root, rank, args.shard_size, resume=mode == "resume")
    seen = valid = 0
    task_counts = Counter()
    started = time.monotonic()
    try:
        for shard_index in owned:
            path = shards[shard_index]
            relative = str(path.relative_to(args.source_root))
            task = path.parents[1].name
            row_index = 0
            for batch in pq.ParquetFile(path).iter_batches(batch_size=args.parquet_batch_size):
                for raw in batch.to_pylist():
                    coordinate = (relative, row_index)
                    if coordinate in cached:
                        row_index += 1
                        continue
                    seen += 1
                    if raw.get("task") != task:
                        raise ValueError(
                            f"task mismatch at {relative}:{row_index}: {raw.get('task')!r}"
                        )
                    source = canonical_source_row(raw)
                    try:
                        item = prepare_row(
                            ar_model=model,
                            processor=processor,
                            config=task,
                            source_shard=relative,
                            source_row=row_index,
                            row=source,
                            max_long_edge=args.max_long_edge,
                            max_sequence_length=args.max_sequence_length,
                            test_fraction=0.0,
                            validation_fraction=0.0,
                            max_vision_patch_tokens=args.max_vision_patch_tokens,
                            image_token_id=model_config.image_token_id,
                            image_preprocessing_profile=args.image_preprocessing_profile,
                            target_canvas_size=args.target_canvas_size,
                            target_background=args.target_background,
                        )
                    except Exception as error:
                        raise RuntimeError(
                            f"invalid ThinkMorph row {relative}:{row_index}"
                        ) from error
                    item["source"]["pid"] = raw["pid"]
                    item["trajectory"]["source"]["dataset"] = "ThinkMorph"
                    item["trajectory"]["source"]["pid"] = raw["pid"]
                    writer.add(item)
                    valid += 1
                    task_counts[task] += 1
                    row_index += 1
                    if seen % args.log_every == 0:
                        elapsed = max(time.monotonic() - started, 1e-6)
                        print(
                            f"PROGRESS rank={rank} seen={seen} valid={valid} "
                            f"rows_per_s={seen / elapsed:.3f} source={relative}",
                            flush=True,
                        )
    finally:
        writer.close()

    write_json(
        args.output_root / "stats" / f"rank-{rank:02d}.json",
        {
            "rank": rank,
            "resume": mode == "resume",
            "cached_before": len(cached),
            "owned_shards": [str(shards[index].relative_to(args.source_root)) for index in owned],
            "seen": seen,
            "valid": valid,
            "tasks": dict(task_counts),
            "elapsed_seconds": time.monotonic() - started,
            "peak_gpu_memory_gib": torch.cuda.max_memory_allocated(device) / 1024**3,
        },
    )
    del model
    torch.cuda.empty_cache()
    dist.barrier()
    if rank == 0:
        stats = _aggregate(args.output_root, world_size, config)
        print("CACHE_COMPLETE " + json.dumps(stats["splits"]), flush=True)
    dist.barrier()
    dist.destroy_process_group()


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--source-root", type=Path, default=Path("data/ThinkMorph"))
    parser.add_argument("--model-path", type=Path, default=Path("models/UniAR-SFT"))
    parser.add_argument("--output-root", type=Path, required=True)
    parser.add_argument("--attn", default="flash_attention_2")
    parser.add_argument("--max-long-edge", type=int, default=960)
    parser.add_argument(
        "--image-preprocessing-profile",
        choices=IMAGE_PREPROCESSING_PROFILES,
        default=LEGACY_IMAGE_PREPROCESSING,
    )
    parser.add_argument("--target-canvas-size", type=int, default=512)
    parser.add_argument("--target-background", type=int, default=128)
    parser.add_argument("--max-sequence-length", type=int, default=32768)
    parser.add_argument("--max-vision-patch-tokens", type=int, default=4096)
    parser.add_argument("--shard-size", type=int, default=64)
    parser.add_argument("--parquet-batch-size", type=int, default=8)
    parser.add_argument("--log-every", type=int, default=100)
    parser.add_argument("--auto-resume", action="store_true")
    return parser


def main() -> None:
    prepare(build_parser().parse_args())


if __name__ == "__main__":
    main()

Xet Storage Details

Size:: 14.8 kB
Xet hash:: 185b7f8bde769aecdd610459cb07076b762bbf160550c0eab9d0c338b18a188e

Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.

	"""Build a train-only canonical UniAR cache from the 24,990 ThinkMorph rows."""

	from __future__ import annotations

	import argparse
	import json
	import os
	import time
	from collections import Counter, defaultdict
	from datetime import timedelta
	from pathlib import Path

	import pyarrow.parquet as pq
	import torch
	import torch.distributed as dist
	from sft.runtime import file_sha256, source_fingerprint, write_json
	from transformers import AutoProcessor
	from uniar import UniARConfig, configure_processor
	from uniar.vision import load_bsq_image_tokenizer
	from zebra_sft.full_data import (
	CACHE_SCHEMA,
	IMAGE_PREPROCESSING_PROFILES,
	LEGACY_IMAGE_PREPROCESSING,
	prepare_row,
	)
	from zebra_sft.prepare_full import CacheWriter
	from zebra_sft.protocol import SERIALIZATION_VERSION, SYSTEM_PROMPT


	TASK_COUNTS = {
	"Chart_Refocus": 6000,
	"Jigsaw_Assembly": 6000,
	"Spatial_Navigation": 6000,
	"Visual_Search": 6990,
	}


	def source_shards(root: Path) -> list[Path]:
	paths = [
	path
	for task in sorted(TASK_COUNTS)
	for path in sorted((root / task / "data").glob("train-*.parquet"))
	]
	if not paths:
	raise FileNotFoundError(f"no ThinkMorph training parquet under {root}")
	return paths


	def canonical_source_row(row: dict) -> dict:
	"""Map one public ThinkMorph row to the canonical typed trajectory inputs."""

	required_text = (
	"pid",
	"question",
	"answer",
	"resoning_thought_0",
	"resoning_thought_1",
	"task",
	)
	for key in required_text:
	if not isinstance(row.get(key), str) or not row[key].strip():
	raise ValueError(f"ThinkMorph row has empty {key}")
	for key in ("problem_image_0", "reasoning_image_0"):
	if not isinstance(row.get(key), dict) or not row[key].get("bytes"):
	raise ValueError(f"ThinkMorph row has no bytes for {key}")
	return {
	"Question": (
	f"{row['question'].strip()}\n"
	"<image_start>[problem_image_0]<image_end>"
	),
	"Text Reasoning Trace": (
	f"{row['resoning_thought_0'].strip()}\n"
	"<image_start>[reasoning_image_0]<image_end>\n"
	f"{row['resoning_thought_1'].strip()}"
	),
	"Final Answer": row["answer"].strip(),
	"problem_image_0": row["problem_image_0"],
	"reasoning_image_0": row["reasoning_image_0"],
	}


	def _mode(output_root: Path, auto_resume: bool) -> str:
	if (output_root / "READY").is_file():
	return "complete"
	if not output_root.exists() or not any(output_root.iterdir()):
	return "fresh"
	if auto_resume:
	return "resume"
	raise FileExistsError(f"cache output already exists: {output_root}")


	def _read_cached_coordinates(manifest_root: Path) -> set[tuple[str, int]]:
	coordinates = set()
	for path in sorted(manifest_root.glob("rank-*.jsonl")):
	data = path.read_bytes()
	boundary = data.rfind(b"\n") + 1
	if boundary != len(data):
	with path.open("r+b") as handle:
	handle.truncate(boundary)
	data = data[:boundary]
	for line in data.decode("utf-8").splitlines():
	if not line:
	continue
	row = json.loads(line)
	coordinate = (row["source"]["shard"], int(row["source"]["row"]))
	if coordinate in coordinates:
	raise RuntimeError(f"duplicate cached coordinate: {coordinate}")
	coordinates.add(coordinate)
	return coordinates


	def _assign_shards(
	shards: list[Path], root: Path, cached: set[tuple[str, int]], world_size: int
	) -> list[list[int]]:
	cached_counts = Counter(shard for shard, _ in cached)
	pending = []
	for index, path in enumerate(shards):
	relative = str(path.relative_to(root))
	remaining = pq.ParquetFile(path).metadata.num_rows - cached_counts[relative]
	if remaining > 0:
	pending.append((remaining, index))
	assignments = [[] for _ in range(world_size)]
	loads = [0] * world_size
	for remaining, index in sorted(pending, reverse=True):
	rank = min(range(world_size), key=lambda value: (loads[value], value))
	assignments[rank].append(index)
	loads[rank] += remaining
	return assignments


	def _aggregate(output_root: Path, world_size: int, config: dict) -> dict:
	rows = []
	rank_stats = []
	for rank in range(world_size):
	manifest = output_root / "manifests" / f"rank-{rank:02d}.jsonl"
	rows.extend(
	json.loads(line)
	for line in manifest.read_text(encoding="utf-8").splitlines()
	if line
	)
	rank_stats.append(
	json.loads(
	(output_root / "stats" / f"rank-{rank:02d}.json").read_text(
	encoding="utf-8"
	)
	)
	)
	rows.sort(key=lambda row: (row["config"], row["source"]["shard"], row["source"]["row"]))
	if len(rows) != sum(TASK_COUNTS.values()):
	raise RuntimeError(f"expected 24990 ThinkMorph rows, got {len(rows)}")
	if len({row["sample_id"] for row in rows}) != len(rows):
	raise RuntimeError("duplicate stable sample IDs in ThinkMorph cache")
	task_counts = Counter(row["config"] for row in rows)
	if dict(task_counts) != TASK_COUNTS:
	raise RuntimeError(f"ThinkMorph task counts differ: {dict(task_counts)}")
	for row in rows:
	row["split"] = "train"

	temporary = output_root / ".manifest.jsonl.tmp"
	with temporary.open("w", encoding="utf-8") as handle:
	for row in rows:
	handle.write(json.dumps(row, ensure_ascii=False) + "\n")
	temporary.replace(output_root / "manifest.jsonl")
	stats = {
	"schema": CACHE_SCHEMA,
	"config": config,
	"samples": len(rows),
	"splits": {"train": len(rows)},
	"config_splits": {
	task: {"train": count} for task, count in sorted(task_counts.items())
	},
	"unique_groups": len({row["group_id"] for row in rows}),
	"group_leakage": 0,
	"selected_cap": dict(Counter(str(row["selected_cap"]) for row in rows)),
	"sequence_length": {
	"max": max(row["sequence_length"] for row in rows),
	"mean": sum(row["sequence_length"] for row in rows) / len(rows),
	},
	"visual_positions": {
	"max": max(row["visual_positions"] for row in rows),
	"mean": sum(row["visual_positions"] for row in rows) / len(rows),
	},
	"rank_stats": rank_stats,
	}
	write_json(output_root / "stats.json", stats)
	ready = output_root / ".READY.tmp"
	ready.write_text("complete\n", encoding="utf-8")
	ready.replace(output_root / "READY")
	return stats


	def prepare(args: argparse.Namespace) -> None:
	mode = _mode(args.output_root, args.auto_resume)
	if mode == "complete":
	print(f"CACHE_ALREADY_COMPLETE {args.output_root}", flush=True)
	return
	if not torch.cuda.is_available():
	raise RuntimeError("ThinkMorph cache preparation requires CUDA")
	local_rank = int(os.environ.get("LOCAL_RANK", 0))
	torch.cuda.set_device(local_rank)
	device = torch.device("cuda", local_rank)
	dist.init_process_group("nccl", device_id=device, timeout=timedelta(hours=6))
	rank, world_size = dist.get_rank(), dist.get_world_size()

	shards = source_shards(args.source_root)
	config = {
	"schema": "thinkmorph_uniar_cache_prepare_v1",
	"source_root": str(args.source_root.resolve()),
	"model_path": str(args.model_path.resolve()),
	"source_fingerprint": source_fingerprint("shared", "sft", "uniar"),
	"producer_sha256": file_sha256(Path(__file__)),
	"full_data_sha256": file_sha256(
	Path(__file__).parents[1] / "zebra_sft" / "full_data.py"
	),
	"trajectory_sha256": file_sha256(
	Path(__file__).parents[1] / "zebra_sft" / "trajectory.py"
	),
	"serialization_version": SERIALIZATION_VERSION,
	"system_prompt": SYSTEM_PROMPT,
	"split_policy": "all_public_rows_are_train",
	"task_counts": TASK_COUNTS,
	"max_long_edge": args.max_long_edge,
	"image_preprocessing_profile": args.image_preprocessing_profile,
	"target_canvas_size": args.target_canvas_size,
	"target_background": args.target_background,
	"max_sequence_length": args.max_sequence_length,
	"shard_size": args.shard_size,
	"world_size": world_size,
	"source_shards": len(shards),
	}
	if rank == 0:
	if mode == "fresh":
	args.output_root.mkdir(parents=True, exist_ok=True)
	(args.output_root / "manifests").mkdir()
	(args.output_root / "stats").mkdir()
	write_json(args.output_root / "config.json", config)
	else:
	previous = json.loads(
	(args.output_root / "config.json").read_text(encoding="utf-8")
	)
	if previous != config:
	raise ValueError("ThinkMorph cache resume configuration mismatch")
	dist.barrier()

	processor = configure_processor(
	AutoProcessor.from_pretrained(args.model_path, use_fast=True)
	)
	model_config = UniARConfig.from_pretrained(args.model_path)
	model = load_bsq_image_tokenizer(
	str(args.model_path),
	subfolder="bsq_encoder",
	attn_implementation=args.attn,
	).to(device).eval()

	cached = (
	_read_cached_coordinates(args.output_root / "manifests")
	if mode == "resume"
	else set()
	)
	assignments = _assign_shards(shards, args.source_root, cached, world_size)
	owned = assignments[rank]
	writer = CacheWriter(args.output_root, rank, args.shard_size, resume=mode == "resume")
	seen = valid = 0
	task_counts = Counter()
	started = time.monotonic()
	try:
	for shard_index in owned:
	path = shards[shard_index]
	relative = str(path.relative_to(args.source_root))
	task = path.parents[1].name
	row_index = 0
	for batch in pq.ParquetFile(path).iter_batches(batch_size=args.parquet_batch_size):
	for raw in batch.to_pylist():
	coordinate = (relative, row_index)
	if coordinate in cached:
	row_index += 1
	continue
	seen += 1
	if raw.get("task") != task:
	raise ValueError(
	f"task mismatch at {relative}:{row_index}: {raw.get('task')!r}"
	)
	source = canonical_source_row(raw)
	try:
	item = prepare_row(
	ar_model=model,
	processor=processor,
	config=task,
	source_shard=relative,
	source_row=row_index,
	row=source,
	max_long_edge=args.max_long_edge,
	max_sequence_length=args.max_sequence_length,
	test_fraction=0.0,
	validation_fraction=0.0,
	max_vision_patch_tokens=args.max_vision_patch_tokens,
	image_token_id=model_config.image_token_id,
	image_preprocessing_profile=args.image_preprocessing_profile,
	target_canvas_size=args.target_canvas_size,
	target_background=args.target_background,
	)
	except Exception as error:
	raise RuntimeError(
	f"invalid ThinkMorph row {relative}:{row_index}"
	) from error
	item["source"]["pid"] = raw["pid"]
	item["trajectory"]["source"]["dataset"] = "ThinkMorph"
	item["trajectory"]["source"]["pid"] = raw["pid"]
	writer.add(item)
	valid += 1
	task_counts[task] += 1
	row_index += 1
	if seen % args.log_every == 0:
	elapsed = max(time.monotonic() - started, 1e-6)
	print(
	f"PROGRESS rank={rank} seen={seen} valid={valid} "
	f"rows_per_s={seen / elapsed:.3f} source={relative}",
	flush=True,
	)
	finally:
	writer.close()

	write_json(
	args.output_root / "stats" / f"rank-{rank:02d}.json",
	{
	"rank": rank,
	"resume": mode == "resume",
	"cached_before": len(cached),
	"owned_shards": [str(shards[index].relative_to(args.source_root)) for index in owned],
	"seen": seen,
	"valid": valid,
	"tasks": dict(task_counts),
	"elapsed_seconds": time.monotonic() - started,
	"peak_gpu_memory_gib": torch.cuda.max_memory_allocated(device) / 1024**3,
	},
	)
	del model
	torch.cuda.empty_cache()
	dist.barrier()
	if rank == 0:
	stats = _aggregate(args.output_root, world_size, config)
	print("CACHE_COMPLETE " + json.dumps(stats["splits"]), flush=True)
	dist.barrier()
	dist.destroy_process_group()


	def build_parser() -> argparse.ArgumentParser:
	parser = argparse.ArgumentParser(description=__doc__)
	parser.add_argument("--source-root", type=Path, default=Path("data/ThinkMorph"))
	parser.add_argument("--model-path", type=Path, default=Path("models/UniAR-SFT"))
	parser.add_argument("--output-root", type=Path, required=True)
	parser.add_argument("--attn", default="flash_attention_2")
	parser.add_argument("--max-long-edge", type=int, default=960)
	parser.add_argument(
	"--image-preprocessing-profile",
	choices=IMAGE_PREPROCESSING_PROFILES,
	default=LEGACY_IMAGE_PREPROCESSING,
	)
	parser.add_argument("--target-canvas-size", type=int, default=512)
	parser.add_argument("--target-background", type=int, default=128)
	parser.add_argument("--max-sequence-length", type=int, default=32768)
	parser.add_argument("--max-vision-patch-tokens", type=int, default=4096)
	parser.add_argument("--shard-size", type=int, default=64)
	parser.add_argument("--parquet-batch-size", type=int, default=8)
	parser.add_argument("--log-every", type=int, default=100)
	parser.add_argument("--auto-resume", action="store_true")
	return parser


	def main() -> None:
	prepare(build_parser().parse_args())


	if __name__ == "__main__":
	main()