Buckets:

EmergentTwilight
/

MMCoT

Files

xet

EmergentTwilight/MMCoT / projects /ThinkMorph-edited /data /convert_thinkmorph.py

EmergentTwilight

19 days ago

download

raw

9.15 kB

"""Convert the four released ThinkMorph datasets to the BAGEL SFT contract."""

from __future__ import annotations

import argparse
import hashlib
import io
import json
import os
import shutil
import uuid
from pathlib import Path

import pyarrow as pa
import pyarrow.parquet as pq
from PIL import Image


TASKS = {
    "Chart_Refocus": ("chart_refocus", 6000, 2),
    "Jigsaw_Assembly": ("jigsaw_assembly", 6000, 3),
    "Spatial_Navigation": ("spatial_navigation", 6000, 1),
    "Visual_Search": ("visual_search", 6990, 3),
}
INTERLEAVED_THINK_SYSTEM_PROMPT = """
Let's think step by step to answer the question. For text-based thinking, enclose the process within <think> </think>, e.g. <think> thinking process here </think>. For visual thinking, enclose the content within <image_start> </image_end>, e.g. <image_start> thinking image here </image_end>. Finally conclude with the final answer wrapped in <answer></answer> tags, i.e.<answer> answer here </answer>.
"""
SOURCE_COLUMNS = (
    "pid",
    "question",
    "answer",
    "problem_image_0",
    "resoning_thought_0",
    "reasoning_image_0",
    "resoning_thought_1",
    "task",
)
OUTPUT_SCHEMA = pa.schema(
    [
        ("image_list", pa.list_(pa.binary())),
        ("instruction_list", pa.list_(pa.string())),
        ("output_text_list", pa.list_(pa.string())),
    ]
)


def _sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(8 * 1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def _image_bytes(value: object, source_file: Path) -> bytes:
    if isinstance(value, dict):
        payload = value.get("bytes")
        if payload is not None:
            return bytes(payload)
        image_path = value.get("path")
        if image_path:
            path = Path(image_path)
            if not path.is_absolute():
                path = source_file.parent / path
            return path.read_bytes()
    if isinstance(value, (bytes, bytearray, memoryview)):
        return bytes(value)
    raise ValueError(f"image payload is missing in {source_file}")


def _image_png(value: object, source_file: Path) -> bytes:
    with Image.open(io.BytesIO(_image_bytes(value, source_file))) as image:
        image.load()
        buffer = io.BytesIO()
        image.save(buffer, format="PNG")
    return buffer.getvalue()


def convert_row(row: dict, source_file: Path) -> dict:
    missing = [name for name in SOURCE_COLUMNS if row.get(name) is None]
    if missing:
        raise ValueError(f"missing {missing} in {source_file}")
    thought_0 = str(row["resoning_thought_0"])
    thought_1 = str(row["resoning_thought_1"])
    answer = str(row["answer"])
    return {
        "image_list": [
            _image_png(row["problem_image_0"], source_file),
            _image_png(row["reasoning_image_0"], source_file),
        ],
        "instruction_list": [
            INTERLEAVED_THINK_SYSTEM_PROMPT + str(row["question"])
        ],
        "output_text_list": [
            f"<think>{thought_0}</think><image_start>",
            f"<image_end><think>{thought_1}</think><answer>{answer}</answer>",
        ],
    }


def _write_chunk(rows: list[dict], path: Path, row_group_size: int) -> None:
    pq.write_table(
        pa.Table.from_pylist(rows, schema=OUTPUT_SCHEMA),
        path,
        compression="zstd",
        row_group_size=row_group_size,
    )


def convert_corpus(
    source_root: Path,
    output_root: Path,
    row_group_size: int,
    rows_per_file: int,
) -> dict:
    if output_root.exists():
        raise FileExistsError(output_root)
    output_root.parent.mkdir(parents=True, exist_ok=True)
    staging = output_root.parent / f".{output_root.name}.tmp-{uuid.uuid4().hex}"
    staging.mkdir()
    source_manifest = []
    output_manifest = []
    parquet_info = {}
    task_counts = {}
    try:
        for task, (slug, expected_rows, expected_files) in TASKS.items():
            sources = sorted((source_root / task / "data").glob("*.parquet"))
            if len(sources) != expected_files:
                raise ValueError(
                    f"{task} has {len(sources)} shards, expected {expected_files}"
                )
            destination_dir = staging / slug
            destination_dir.mkdir()
            task_count = 0
            output_count = 0
            pending = []
            for source in sources:
                parquet = pq.ParquetFile(source)
                source_manifest.append(
                    {
                        "task": task,
                        "path": str(source),
                        "sha256": _sha256(source),
                        "rows": parquet.metadata.num_rows,
                    }
                )
                for batch in parquet.iter_batches(
                    batch_size=rows_per_file,
                    columns=list(SOURCE_COLUMNS),
                ):
                    rows = batch.to_pylist()
                    mismatched = [row["task"] for row in rows if row["task"] != task]
                    if mismatched:
                        raise ValueError(f"task mismatch in {source}: {mismatched[0]}")
                    pending.extend(convert_row(row, source) for row in rows)
                    while len(pending) >= rows_per_file:
                        chunk, pending = pending[:rows_per_file], pending[rows_per_file:]
                        destination = (
                            destination_dir / f"chunk_{output_count:04d}.parquet"
                        )
                        _write_chunk(chunk, destination, row_group_size)
                        output_count += 1
                        task_count += len(chunk)
            if pending:
                destination = destination_dir / f"chunk_{output_count:04d}.parquet"
                _write_chunk(pending, destination, row_group_size)
                output_count += 1
                task_count += len(pending)
            if task_count != expected_rows:
                raise ValueError(f"{task} has {task_count} rows, expected {expected_rows}")
            task_counts[task] = task_count
            for destination in sorted(destination_dir.glob("*.parquet")):
                parquet = pq.ParquetFile(destination)
                final_path = output_root / slug / destination.name
                parquet_info[str(final_path)] = {
                    "num_row_groups": parquet.num_row_groups,
                    "num_rows": parquet.metadata.num_rows,
                }
                output_manifest.append(
                    {
                        "task": task,
                        "path": str(final_path),
                        "sha256": _sha256(destination),
                        "rows": parquet.metadata.num_rows,
                    }
                )
            print(
                f"CONVERT_TASK_COMPLETE task={task} rows={task_count} "
                f"files={output_count}",
                flush=True,
            )
        (staging / "parquet_info.json").write_text(
            json.dumps(parquet_info, indent=2, sort_keys=True) + "\n",
            encoding="utf-8",
        )
        summary = {
            "schema": "thinkmorph_bagel_sft_v1",
            "source_root": str(source_root),
            "rows": sum(task_counts.values()),
            "tasks": task_counts,
            "row_group_size": row_group_size,
            "rows_per_file": rows_per_file,
            "source_files": source_manifest,
            "output_files": output_manifest,
        }
        (staging / "manifest.json").write_text(
            json.dumps(summary, indent=2, sort_keys=True) + "\n",
            encoding="utf-8",
        )
        (staging / "READY").write_text("complete\n", encoding="utf-8")
        os.replace(staging, output_root)
        return summary
    except BaseException:
        shutil.rmtree(staging, ignore_errors=True)
        raise


def main() -> None:
    project_root = Path(
        os.environ.get(
            "SPATIAL_MMCOT_PROJECT_ROOT",
            Path(__file__).resolve().parents[3],
        )
    )
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--source-root",
        type=Path,
        default=project_root / "data/ThinkMorph",
    )
    parser.add_argument(
        "--output-root",
        type=Path,
        default=project_root / "data/ThinkMorph/Bagel-SFT",
    )
    parser.add_argument("--row-group-size", type=int, default=10)
    parser.add_argument("--rows-per-file", type=int, default=100)
    args = parser.parse_args()
    if (
        args.row_group_size <= 0
        or args.rows_per_file <= 0
        or args.rows_per_file % args.row_group_size
    ):
        raise ValueError("rows-per-file must be a positive multiple of row-group-size")
    summary = convert_corpus(
        args.source_root.resolve(),
        args.output_root.resolve(),
        args.row_group_size,
        args.rows_per_file,
    )
    print(json.dumps({key: value for key, value in summary.items() if not key.endswith("_files")}))


if __name__ == "__main__":
    main()

Xet Storage Details

Size:: 9.15 kB
Xet hash:: e4b3b31772b6836651e56863d0230bc7ee185be2250884c2f81aea6bf774fbdc

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

	"""Convert the four released ThinkMorph datasets to the BAGEL SFT contract."""

	from __future__ import annotations

	import argparse
	import hashlib
	import io
	import json
	import os
	import shutil
	import uuid
	from pathlib import Path

	import pyarrow as pa
	import pyarrow.parquet as pq
	from PIL import Image


	TASKS = {
	"Chart_Refocus": ("chart_refocus", 6000, 2),
	"Jigsaw_Assembly": ("jigsaw_assembly", 6000, 3),
	"Spatial_Navigation": ("spatial_navigation", 6000, 1),
	"Visual_Search": ("visual_search", 6990, 3),
	}
	INTERLEAVED_THINK_SYSTEM_PROMPT = """
	Let's think step by step to answer the question. For text-based thinking, enclose the process within <think> </think>, e.g. <think> thinking process here </think>. For visual thinking, enclose the content within <image_start> </image_end>, e.g. <image_start> thinking image here </image_end>. Finally conclude with the final answer wrapped in <answer></answer> tags, i.e.<answer> answer here </answer>.
	"""
	SOURCE_COLUMNS = (
	"pid",
	"question",
	"answer",
	"problem_image_0",
	"resoning_thought_0",
	"reasoning_image_0",
	"resoning_thought_1",
	"task",
	)
	OUTPUT_SCHEMA = pa.schema(
	[
	("image_list", pa.list_(pa.binary())),
	("instruction_list", pa.list_(pa.string())),
	("output_text_list", pa.list_(pa.string())),
	]
	)


	def _sha256(path: Path) -> str:
	digest = hashlib.sha256()
	with path.open("rb") as handle:
	for chunk in iter(lambda: handle.read(8 * 1024 * 1024), b""):
	digest.update(chunk)
	return digest.hexdigest()


	def _image_bytes(value: object, source_file: Path) -> bytes:
	if isinstance(value, dict):
	payload = value.get("bytes")
	if payload is not None:
	return bytes(payload)
	image_path = value.get("path")
	if image_path:
	path = Path(image_path)
	if not path.is_absolute():
	path = source_file.parent / path
	return path.read_bytes()
	if isinstance(value, (bytes, bytearray, memoryview)):
	return bytes(value)
	raise ValueError(f"image payload is missing in {source_file}")


	def _image_png(value: object, source_file: Path) -> bytes:
	with Image.open(io.BytesIO(_image_bytes(value, source_file))) as image:
	image.load()
	buffer = io.BytesIO()
	image.save(buffer, format="PNG")
	return buffer.getvalue()


	def convert_row(row: dict, source_file: Path) -> dict:
	missing = [name for name in SOURCE_COLUMNS if row.get(name) is None]
	if missing:
	raise ValueError(f"missing {missing} in {source_file}")
	thought_0 = str(row["resoning_thought_0"])
	thought_1 = str(row["resoning_thought_1"])
	answer = str(row["answer"])
	return {
	"image_list": [
	_image_png(row["problem_image_0"], source_file),
	_image_png(row["reasoning_image_0"], source_file),
	],
	"instruction_list": [
	INTERLEAVED_THINK_SYSTEM_PROMPT + str(row["question"])
	],
	"output_text_list": [
	f"<think>{thought_0}</think><image_start>",
	f"<image_end><think>{thought_1}</think><answer>{answer}</answer>",
	],
	}


	def _write_chunk(rows: list[dict], path: Path, row_group_size: int) -> None:
	pq.write_table(
	pa.Table.from_pylist(rows, schema=OUTPUT_SCHEMA),
	path,
	compression="zstd",
	row_group_size=row_group_size,
	)


	def convert_corpus(
	source_root: Path,
	output_root: Path,
	row_group_size: int,
	rows_per_file: int,
	) -> dict:
	if output_root.exists():
	raise FileExistsError(output_root)
	output_root.parent.mkdir(parents=True, exist_ok=True)
	staging = output_root.parent / f".{output_root.name}.tmp-{uuid.uuid4().hex}"
	staging.mkdir()
	source_manifest = []
	output_manifest = []
	parquet_info = {}
	task_counts = {}
	try:
	for task, (slug, expected_rows, expected_files) in TASKS.items():
	sources = sorted((source_root / task / "data").glob("*.parquet"))
	if len(sources) != expected_files:
	raise ValueError(
	f"{task} has {len(sources)} shards, expected {expected_files}"
	)
	destination_dir = staging / slug
	destination_dir.mkdir()
	task_count = 0
	output_count = 0
	pending = []
	for source in sources:
	parquet = pq.ParquetFile(source)
	source_manifest.append(
	{
	"task": task,
	"path": str(source),
	"sha256": _sha256(source),
	"rows": parquet.metadata.num_rows,
	}
	)
	for batch in parquet.iter_batches(
	batch_size=rows_per_file,
	columns=list(SOURCE_COLUMNS),
	):
	rows = batch.to_pylist()
	mismatched = [row["task"] for row in rows if row["task"] != task]
	if mismatched:
	raise ValueError(f"task mismatch in {source}: {mismatched[0]}")
	pending.extend(convert_row(row, source) for row in rows)
	while len(pending) >= rows_per_file:
	chunk, pending = pending[:rows_per_file], pending[rows_per_file:]
	destination = (
	destination_dir / f"chunk_{output_count:04d}.parquet"
	)
	_write_chunk(chunk, destination, row_group_size)
	output_count += 1
	task_count += len(chunk)
	if pending:
	destination = destination_dir / f"chunk_{output_count:04d}.parquet"
	_write_chunk(pending, destination, row_group_size)
	output_count += 1
	task_count += len(pending)
	if task_count != expected_rows:
	raise ValueError(f"{task} has {task_count} rows, expected {expected_rows}")
	task_counts[task] = task_count
	for destination in sorted(destination_dir.glob("*.parquet")):
	parquet = pq.ParquetFile(destination)
	final_path = output_root / slug / destination.name
	parquet_info[str(final_path)] = {
	"num_row_groups": parquet.num_row_groups,
	"num_rows": parquet.metadata.num_rows,
	}
	output_manifest.append(
	{
	"task": task,
	"path": str(final_path),
	"sha256": _sha256(destination),
	"rows": parquet.metadata.num_rows,
	}
	)
	print(
	f"CONVERT_TASK_COMPLETE task={task} rows={task_count} "
	f"files={output_count}",
	flush=True,
	)
	(staging / "parquet_info.json").write_text(
	json.dumps(parquet_info, indent=2, sort_keys=True) + "\n",
	encoding="utf-8",
	)
	summary = {
	"schema": "thinkmorph_bagel_sft_v1",
	"source_root": str(source_root),
	"rows": sum(task_counts.values()),
	"tasks": task_counts,
	"row_group_size": row_group_size,
	"rows_per_file": rows_per_file,
	"source_files": source_manifest,
	"output_files": output_manifest,
	}
	(staging / "manifest.json").write_text(
	json.dumps(summary, indent=2, sort_keys=True) + "\n",
	encoding="utf-8",
	)
	(staging / "READY").write_text("complete\n", encoding="utf-8")
	os.replace(staging, output_root)
	return summary
	except BaseException:
	shutil.rmtree(staging, ignore_errors=True)
	raise


	def main() -> None:
	project_root = Path(
	os.environ.get(
	"SPATIAL_MMCOT_PROJECT_ROOT",
	Path(__file__).resolve().parents[3],
	)
	)
	parser = argparse.ArgumentParser(description=__doc__)
	parser.add_argument(
	"--source-root",
	type=Path,
	default=project_root / "data/ThinkMorph",
	)
	parser.add_argument(
	"--output-root",
	type=Path,
	default=project_root / "data/ThinkMorph/Bagel-SFT",
	)
	parser.add_argument("--row-group-size", type=int, default=10)
	parser.add_argument("--rows-per-file", type=int, default=100)
	args = parser.parse_args()
	if (
	args.row_group_size <= 0
	or args.rows_per_file <= 0
	or args.rows_per_file % args.row_group_size
	):
	raise ValueError("rows-per-file must be a positive multiple of row-group-size")
	summary = convert_corpus(
	args.source_root.resolve(),
	args.output_root.resolve(),
	args.row_group_size,
	args.rows_per_file,
	)
	print(json.dumps({key: value for key, value in summary.items() if not key.endswith("_files")}))


	if __name__ == "__main__":
	main()