rs / satlink /train /data /make_verl_opsd_data.py

satlink/ at top level, with train/ (OPSD + data build) included

0eb0c3a verified 2 months ago

6.88 kB

"""Build the verl / Vision-OPD parquet for on-policy self-distillation on our reconstructions.
  student `images`         = 48-token-per-crop One-D-Piece reconstruction
  teacher `teacher_images` = the ORIGINAL image
Vision-OPD's teacher path (`_swap_images_in_messages`) reuses the student's prompt verbatim and
only swaps the pixels, which is exactly the comparison we want: same question, privileged input.
Images stay as embedded bytes rather than paths -- verl's `process_image` accepts
`{"bytes": ...}`, so the 382k reconstructions never have to be written out as loose files again.
The two source parquet sets were generated from the same manifest in the same order and verified
row-aligned (identical `messages`), so they are zipped by position.
  python -m mllm.make_verl_opsd_data --out data/verl_opsd --limit 2000     # smoke
  python -m mllm.make_verl_opsd_data --out data/verl_opsd                  # full
"""
import os, sys, glob, json, argparse

import os, sys
_ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
for _p in (_ROOT, os.path.join(_ROOT, "src")):
    if _p not in sys.path:
        sys.path.insert(0, _p)

ROOT = "${TOKLIP_ROOT}"
RECON = f"{ROOT}/data/GeoCoT380k_parquet_v2"   # native-resolution overlapping tiles
ORIG = f"{ROOT}/data/GeoCoT380k_orig_parquet"


def _match_geometry(orig_bytes, recon_bytes):
    """Re-encode the original at the reconstruction's exact size, so the ONLY difference between
    the teacher's and the student's view is the 48-token bottleneck."""
    import io as _io
    from PIL import Image
    rec = Image.open(_io.BytesIO(recon_bytes))
    im = Image.open(_io.BytesIO(orig_bytes))
    if im.size == rec.size:
        # Native-resolution tiling makes the reconstruction the SAME size as its source, so this
        # is the common case. Returning the bytes untouched skips a full JPEG decode+re-encode
        # per row -- that round trip was the entire cost of this pass.
        return orig_bytes
    b = _io.BytesIO()
    im.convert("RGB").resize(rec.size, Image.BICUBIC).save(b, format="JPEG", quality=95)
    return b.getvalue()


def rows(d, files=None):
    import pyarrow.parquet as pq
    for f in (files or sorted(glob.glob(f"{d}/*.parquet"))):
        for b in pq.ParquetFile(f).iter_batches(batch_size=64):
            yield from b.to_pylist()


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--out", required=True)
    ap.add_argument("--limit", type=int, default=0)
    ap.add_argument("--shard-size", type=int, default=20000)
    ap.add_argument("--data-source", default="geocot380k_recon_opsd")
    ap.add_argument("--recon", default=RECON)
    ap.add_argument("--rank", type=int, default=0)
    ap.add_argument("--workers", type=int, default=1)
    a = ap.parse_args()

    import pyarrow.parquet as pq
    from datasets import Dataset, Features, Sequence, Value

    os.makedirs(a.out, exist_ok=True)
    buf, shard, n, skipped = [], 0, 0, 0

    # Write through `datasets`, NOT pq.write_table. Raw pyarrow lays the nested
    # list<struct<bytes,path>> out as a chunked array, and verl's datasets.load_dataset then dies
    # with "Nested data conversions not implemented for chunked array outputs" -- the same trap
    # the reconstruction packer already had to work around.
    img = {"bytes": Value("binary"), "path": Value("string")}
    feats = Features({
        "data_source": Value("string"),
        "prompt": [{"role": Value("string"), "content": Value("string")}],
        "images": [img],
        "teacher_images": [img],
        "ability": Value("string"),
        "reward_model": {"style": Value("string"), "ground_truth": Value("string")},
        "extra_info": {"answer": Value("string"), "question": Value("string")},
    })

    def flush():
        nonlocal buf, shard
        if not buf:
            return
        tmp = f"{a.out}/opsd-{shard:05d}.parquet.tmp"
        Dataset.from_list(buf, features=feats).to_parquet(tmp)
        next(pq.ParquetFile(tmp).iter_batches(batch_size=4))      # readable before publishing
        os.replace(tmp, f"{a.out}/opsd-{shard:05d}.parquet")
        print(f"  wrote opsd-{shard:05d}.parquet ({len(buf)} rows)", flush=True)
        buf, shard = [], shard + a.workers

    # One source shard per worker: recon and orig shards share names and row order, so slicing by
    # file keeps every (student, teacher) pair together.
    rf = sorted(glob.glob(f"{a.recon}/*.parquet"))
    of = sorted(glob.glob(f"{ORIG}/*.parquet"))
    mine = [(x, y) for i, (x, y) in enumerate(zip(rf, of)) if i % a.workers == a.rank]
    shard = a.rank
    import time as _t
    t0 = _t.time()
    for r, o in zip(rows(None, [x for x, _ in mine]), rows(None, [y for _, y in mine])):
        # positional zip is only safe because the two sets are row-aligned; assert it per row
        # rather than trusting the earlier spot check -- a silent misalignment would pair every
        # student image with the WRONG original and quietly destroy the distillation signal.
        if r["messages"][0]["content"] != o["messages"][0]["content"]:
            skipped += 1
            continue
        rb = (r["images"] or [{}])[0].get("bytes")
        ob = (o["images"] or [{}])[0].get("bytes")
        if not rb or not ob:
            skipped += 1
            continue
        prompt = r["messages"][0]["content"]
        answer = r["messages"][1]["content"]
        # Put the ORIGINAL at the RECONSTRUCTION's geometry. Left native, a 512x512 source would
        # reach the teacher at 512 while the student sees the 768 reconstruction, and grounding is
        # extremely resolution-sensitive here (same model, same image: 2.5% @0.5 at native 512 vs
        # 16.7% upscaled to 768). The teacher would then hold an advantage the student cannot ever
        # close, and the distillation target would stop being "what the 48 tokens lost".
        ob = _match_geometry(ob, rb)
        buf.append({
            "data_source": a.data_source,
            "prompt": [{"role": "user", "content": prompt}],
            "images": [{"bytes": rb, "path": None}],            # student: reconstruction
            "teacher_images": [{"bytes": ob, "path": None}],    # teacher: original, same geometry
            "ability": "remote_sensing_cot",
            "reward_model": {"style": "none", "ground_truth": answer},
            "extra_info": {"answer": answer,
                           "question": prompt.replace("<image>", "").strip()},
        })
        n += 1
        if n % 2000 == 0:
            print(f"  [rank {a.rank}] {n} rows | {n/(_t.time()-t0):.0f} rows/s", flush=True)
        if len(buf) >= a.shard_size:
            flush()
        if a.limit and n >= a.limit:
            break
    flush()
    print(f"DONE: {n} rows in {shard} shards -> {a.out}   (skipped {skipped})")


if __name__ == "__main__":
    main()

	"""Build the verl / Vision-OPD parquet for on-policy self-distillation on our reconstructions.

	student `images` = 48-token-per-crop One-D-Piece reconstruction
	teacher `teacher_images` = the ORIGINAL image

	Vision-OPD's teacher path (`_swap_images_in_messages`) reuses the student's prompt verbatim and
	only swaps the pixels, which is exactly the comparison we want: same question, privileged input.

	Images stay as embedded bytes rather than paths -- verl's `process_image` accepts
	`{"bytes": ...}`, so the 382k reconstructions never have to be written out as loose files again.

	The two source parquet sets were generated from the same manifest in the same order and verified
	row-aligned (identical `messages`), so they are zipped by position.

	python -m mllm.make_verl_opsd_data --out data/verl_opsd --limit 2000 # smoke
	python -m mllm.make_verl_opsd_data --out data/verl_opsd # full
	"""
	import os, sys, glob, json, argparse

	import os, sys
	_ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
	for _p in (_ROOT, os.path.join(_ROOT, "src")):
	if _p not in sys.path:
	sys.path.insert(0, _p)

	ROOT = "${TOKLIP_ROOT}"
	RECON = f"{ROOT}/data/GeoCoT380k_parquet_v2" # native-resolution overlapping tiles
	ORIG = f"{ROOT}/data/GeoCoT380k_orig_parquet"


	def _match_geometry(orig_bytes, recon_bytes):
	"""Re-encode the original at the reconstruction's exact size, so the ONLY difference between
	the teacher's and the student's view is the 48-token bottleneck."""
	import io as _io
	from PIL import Image
	rec = Image.open(_io.BytesIO(recon_bytes))
	im = Image.open(_io.BytesIO(orig_bytes))
	if im.size == rec.size:
	# Native-resolution tiling makes the reconstruction the SAME size as its source, so this
	# is the common case. Returning the bytes untouched skips a full JPEG decode+re-encode
	# per row -- that round trip was the entire cost of this pass.
	return orig_bytes
	b = _io.BytesIO()
	im.convert("RGB").resize(rec.size, Image.BICUBIC).save(b, format="JPEG", quality=95)
	return b.getvalue()


	def rows(d, files=None):
	import pyarrow.parquet as pq
	for f in (files or sorted(glob.glob(f"{d}/*.parquet"))):
	for b in pq.ParquetFile(f).iter_batches(batch_size=64):
	yield from b.to_pylist()


	def main():
	ap = argparse.ArgumentParser()
	ap.add_argument("--out", required=True)
	ap.add_argument("--limit", type=int, default=0)
	ap.add_argument("--shard-size", type=int, default=20000)
	ap.add_argument("--data-source", default="geocot380k_recon_opsd")
	ap.add_argument("--recon", default=RECON)
	ap.add_argument("--rank", type=int, default=0)
	ap.add_argument("--workers", type=int, default=1)
	a = ap.parse_args()

	import pyarrow.parquet as pq
	from datasets import Dataset, Features, Sequence, Value

	os.makedirs(a.out, exist_ok=True)
	buf, shard, n, skipped = [], 0, 0, 0

	# Write through `datasets`, NOT pq.write_table. Raw pyarrow lays the nested
	# list<struct<bytes,path>> out as a chunked array, and verl's datasets.load_dataset then dies
	# with "Nested data conversions not implemented for chunked array outputs" -- the same trap
	# the reconstruction packer already had to work around.
	img = {"bytes": Value("binary"), "path": Value("string")}
	feats = Features({
	"data_source": Value("string"),
	"prompt": [{"role": Value("string"), "content": Value("string")}],
	"images": [img],
	"teacher_images": [img],
	"ability": Value("string"),
	"reward_model": {"style": Value("string"), "ground_truth": Value("string")},
	"extra_info": {"answer": Value("string"), "question": Value("string")},
	})

	def flush():
	nonlocal buf, shard
	if not buf:
	return
	tmp = f"{a.out}/opsd-{shard:05d}.parquet.tmp"
	Dataset.from_list(buf, features=feats).to_parquet(tmp)
	next(pq.ParquetFile(tmp).iter_batches(batch_size=4)) # readable before publishing
	os.replace(tmp, f"{a.out}/opsd-{shard:05d}.parquet")
	print(f" wrote opsd-{shard:05d}.parquet ({len(buf)} rows)", flush=True)
	buf, shard = [], shard + a.workers

	# One source shard per worker: recon and orig shards share names and row order, so slicing by
	# file keeps every (student, teacher) pair together.
	rf = sorted(glob.glob(f"{a.recon}/*.parquet"))
	of = sorted(glob.glob(f"{ORIG}/*.parquet"))
	mine = [(x, y) for i, (x, y) in enumerate(zip(rf, of)) if i % a.workers == a.rank]
	shard = a.rank
	import time as _t
	t0 = _t.time()
	for r, o in zip(rows(None, [x for x, _ in mine]), rows(None, [y for _, y in mine])):
	# positional zip is only safe because the two sets are row-aligned; assert it per row
	# rather than trusting the earlier spot check -- a silent misalignment would pair every
	# student image with the WRONG original and quietly destroy the distillation signal.
	if r["messages"][0]["content"] != o["messages"][0]["content"]:
	skipped += 1
	continue
	rb = (r["images"] or [{}])[0].get("bytes")
	ob = (o["images"] or [{}])[0].get("bytes")
	if not rb or not ob:
	skipped += 1
	continue
	prompt = r["messages"][0]["content"]
	answer = r["messages"][1]["content"]
	# Put the ORIGINAL at the RECONSTRUCTION's geometry. Left native, a 512x512 source would
	# reach the teacher at 512 while the student sees the 768 reconstruction, and grounding is
	# extremely resolution-sensitive here (same model, same image: 2.5% @0.5 at native 512 vs
	# 16.7% upscaled to 768). The teacher would then hold an advantage the student cannot ever
	# close, and the distillation target would stop being "what the 48 tokens lost".
	ob = _match_geometry(ob, rb)
	buf.append({
	"data_source": a.data_source,
	"prompt": [{"role": "user", "content": prompt}],
	"images": [{"bytes": rb, "path": None}], # student: reconstruction
	"teacher_images": [{"bytes": ob, "path": None}], # teacher: original, same geometry
	"ability": "remote_sensing_cot",
	"reward_model": {"style": "none", "ground_truth": answer},
	"extra_info": {"answer": answer,
	"question": prompt.replace("<image>", "").strip()},
	})
	n += 1
	if n % 2000 == 0:
	print(f" [rank {a.rank}] {n} rows \| {n/(_t.time()-t0):.0f} rows/s", flush=True)
	if len(buf) >= a.shard_size:
	flush()
	if a.limit and n >= a.limit:
	break
	flush()
	print(f"DONE: {n} rows in {shard} shards -> {a.out} (skipped {skipped})")


	if __name__ == "__main__":
	main()