[RL] thinker video+image input

examples/omni_thinker/bench_mm_payload.py

@@ -0,0 +1,132 @@
+"""Benchmark the processed-multimodal wire contract (serialize -> JSON -> deserialize).
+
+Quantifies the base64-tensor-in-JSON costs of `multimodal_train_inputs` per modality —
+the known heavy case is video (pixel_values_videos grows with frames x resolution) — so
+the transport decision (keep JSON vs move to a binary/shared-memory path) is made on
+numbers, not vibes.
+
+    python examples/omni_thinker/bench_mm_payload.py [--frames 8 16 32] [--seconds 10 30]
+    python examples/omni_thinker/bench_mm_payload.py --server http://<ip>:<port>  # + HTTP POST timing
+"""
+
+import argparse
+import json
+import time
+
+import torch
+
+from miles.rollout.generate_utils.generate_endpoint_utils import serialize_multimodal_train_inputs
+
+PATCH = 16
+TEMPORAL_PATCH = 2
+MERGE = 2  # spatial merge; grid counts are pre-merge patches per qwen vision geometry
+MEL_BINS = 128
+FRAMES_PER_SECOND_MEL = 100  # whisper-style feature extractor: 10ms hop
+
+
+def make_video_bundle(num_frames: int, height: int, width: int) -> dict[str, torch.Tensor]:
+    grid_t = num_frames // TEMPORAL_PATCH
+    grid_h, grid_w = height // PATCH, width // PATCH
+    patch_dim = 3 * TEMPORAL_PATCH * PATCH * PATCH
+    return {
+        "pixel_values_videos": torch.randn(grid_t * grid_h * grid_w, patch_dim, dtype=torch.float32),
+        "video_grid_thw": torch.tensor([[grid_t, grid_h, grid_w]], dtype=torch.long),
+        "video_second_per_grid": torch.tensor([0.5], dtype=torch.float32),
+    }
+
+
+def make_image_bundle(height: int, width: int) -> dict[str, torch.Tensor]:
+    grid_h, grid_w = height // PATCH, width // PATCH
+    patch_dim = 3 * TEMPORAL_PATCH * PATCH * PATCH
+    return {
+        "pixel_values": torch.randn(grid_h * grid_w, patch_dim, dtype=torch.float32),
+        "image_grid_thw": torch.tensor([[1, grid_h, grid_w]], dtype=torch.long),
+    }
+
+
+def make_audio_bundle(seconds: float) -> dict[str, torch.Tensor]:
+    frames = int(seconds * FRAMES_PER_SECOND_MEL)
+    return {
+        "input_features": torch.randn(1, MEL_BINS, frames, dtype=torch.float32),
+        "feature_attention_mask": torch.ones(1, frames, dtype=torch.long),
+    }
+
+
+def _deserialize(bundle: dict) -> dict[str, torch.Tensor]:
+    """Mirror the server decode (sglang-omni preprocessor.py)."""
+    import pybase64
+
+    out = {}
+    for name, spec in bundle["tensors"].items():
+        raw = bytearray(pybase64.b64decode(spec["data"]))
+        out[name] = torch.frombuffer(raw, dtype=getattr(torch, spec["dtype"])).reshape(spec["shape"])
+    return out
+
+
+def measure_roundtrip(tensors: dict[str, torch.Tensor], server: str | None = None) -> dict:
+    t0 = time.perf_counter()
+    bundle = serialize_multimodal_train_inputs(tensors)
+    t1 = time.perf_counter()
+    payload = json.dumps(bundle)
+    t2 = time.perf_counter()
+    decoded_bundle = json.loads(payload)
+    restored = _deserialize(decoded_bundle)
+    t3 = time.perf_counter()
+
+    report = {
+        "raw_bytes": sum(t.numel() * t.element_size() for t in tensors.values()),
+        "payload_bytes": len(payload),
+        "serialize_ms": (t1 - t0) * 1e3,
+        "json_dump_ms": (t2 - t1) * 1e3,
+        "deserialize_ms": (t3 - t2) * 1e3,
+        "roundtrip_equal": all(torch.equal(restored[k], tensors[k].cpu()) for k in tensors),
+    }
+
+    if server:
+        import requests
+
+        body = {
+            "input_ids": [1, 2, 3],
+            "sampling_params": {"max_new_tokens": 1, "temperature": 0.0},
+            "multimodal_train_inputs": bundle,
+        }
+        t4 = time.perf_counter()
+        resp = requests.post(f"{server}/generate", json=body, timeout=600)
+        report["http_ms"] = (time.perf_counter() - t4) * 1e3
+        report["http_status"] = resp.status_code
+    return report
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--frames", type=int, nargs="+", default=[8, 16, 32, 64])
+    parser.add_argument("--resolution", type=int, default=448)
+    parser.add_argument("--seconds", type=float, nargs="+", default=[10.0, 30.0])
+    parser.add_argument("--server", type=str, default=None, help="optional http://ip:port for POST timing")
+    args = parser.parse_args()
+
+    rows = [("image", f"{args.resolution}px", make_image_bundle(args.resolution, args.resolution))]
+    rows += [
+        ("video", f"{n}f@{args.resolution}px", make_video_bundle(n, args.resolution, args.resolution))
+        for n in args.frames
+    ]
+    rows += [("audio", f"{s:.0f}s", make_audio_bundle(s)) for s in args.seconds]
+
+    header = (
+        f"{'modality':8s} {'case':12s} {'raw_MB':>8s} {'json_MB':>8s} {'ser_ms':>8s} {'dump_ms':>8s} {'deser_ms':>9s}"
+    )
+    print(header)
+    for modality, case, tensors in rows:
+        r = measure_roundtrip(tensors, server=args.server)
+        line = (
+            f"{modality:8s} {case:12s} {r['raw_bytes'] / 2**20:8.1f} {r['payload_bytes'] / 2**20:8.1f} "
+            f"{r['serialize_ms']:8.1f} {r['json_dump_ms']:8.1f} {r['deserialize_ms']:9.1f}"
+        )
+        if args.server:
+            line += f"  http={r['http_ms']:.0f}ms({r['http_status']})"
+        assert r["roundtrip_equal"]
+        print(line)
+
+
+if __name__ == "__main__":
+    main()

examples/omni_thinker/prepare_avqa.py

@@ -0,0 +1,67 @@
+"""Convert AVQA (Joysw909/AVQA, r1aqa line-json) into a miles jsonl for audio-input GRPO.
+
+Each output row is the miles multimodal convention: a `<audio>` placeholder prompt (paired
+with --multimodal-keys '{"audio": "audios"}'), the option letter as label, and the raw
+choices in metadata so `--rm-type gpqa` scores letter or option-text answers.
+
+    python examples/omni_thinker/prepare_avqa.py \
+        --src <dataset_dir>/train_r1aqa_line.json --dst <data_dir>/avqa.jsonl \
+        --audio-root <dataset_dir> [--max-samples 5000]
+"""
+
+import argparse
+import json
+import string
+from pathlib import Path
+
+# AVQA questions say "in the video", but the r1aqa recipe feeds the audio track only; the
+# questions are answerable from sound (VGGSound clips).
+_PROMPT_TEMPLATE = (
+    "<audio>{question}\n{options}\n"
+    "Listen to the audio and choose the best option. Respond with the option letter in the form 'Answer: <letter>'."
+)
+_DATASET_PATH_PREFIX = "./Joysw909/AVQA/"
+
+
+def convert_row(row: dict, audio_root: str) -> dict:
+    choices = list(row["multi_choice"])
+    answer_index = int(row["answer"])
+    assert 0 <= answer_index < len(choices) <= len(string.ascii_uppercase), f"bad row: {row}"
+    letters = string.ascii_uppercase[: len(choices)]
+    options = "\n".join(f"{letter}. {choice}" for letter, choice in zip(letters, choices, strict=True))
+    audio_rel = row["audio_path"].removeprefix(_DATASET_PATH_PREFIX)
+    return {
+        "prompt": _PROMPT_TEMPLATE.format(question=row["question_text"].strip(), options=options),
+        "audios": [str(Path(audio_root) / audio_rel)],
+        "label": letters[answer_index],
+        "metadata": {"choices": choices},
+    }
+
+
+def convert_file(src, dst, audio_root: str, max_samples: int | None = None) -> int:
+    n = 0
+    with open(src) as fin, open(dst, "w") as fout:
+        for line in fin:
+            if not line.strip():
+                continue
+            if max_samples is not None and n >= max_samples:
+                break
+            fout.write(json.dumps(convert_row(json.loads(line), audio_root)) + "\n")
+            n += 1
+    assert n > 0, f"no rows converted from {src}"
+    return n
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--src", required=True, help="train_r1aqa_line.json from Joysw909/AVQA")
+    parser.add_argument("--dst", required=True, help="output jsonl path")
+    parser.add_argument("--audio-root", required=True, help="dataset dir containing VGG*/*.wav")
+    parser.add_argument("--max-samples", type=int, default=None)
+    args = parser.parse_args()
+    n = convert_file(args.src, args.dst, audio_root=args.audio_root, max_samples=args.max_samples)
+    print(f"[done] {n} samples -> {args.dst}")
+
+
+if __name__ == "__main__":
+    main()

examples/omni_thinker/prepare_video_r1.py

@@ -0,0 +1,99 @@
+"""Convert Video-R1-260k (mixed image/video MCQ with verifiable answers) into miles jsonl.
+
+Rows carry lettered options and `<answer>X</answer>` solutions, so `--rm-type gpqa`
+scores them directly. Both media columns are always present (one empty), pairing with
+`--multimodal-keys '{"video": "videos", "image": "images"}'`.
+
+    python examples/omni_thinker/prepare_video_r1.py \
+        --src <dir>/Video-R1-260k.json --dst <data_dir>/video_r1.jsonl \
+        --media-root <dir> [--data-types video image] [--path-prefixes CLEVRER] [--max-samples 5000]
+"""
+
+import argparse
+import json
+import re
+from pathlib import Path
+
+_PLACEHOLDER = {"video": "<video>", "image": "<image>"}
+_PROMPT_TEMPLATE = (
+    "{placeholder}{problem}\n{options}\n"
+    "Watch carefully and choose the best option. Respond with the option letter in the form 'Answer: <letter>'."
+)
+_ANSWER_RE = re.compile(r"<answer>\s*([A-Za-z])\s*</answer>")
+
+
+def convert_row(row: dict, media_root) -> dict:
+    match = _ANSWER_RE.search(row["solution"])
+    assert match, f"no letter answer in solution: {row['solution']!r} (problem_id={row.get('problem_id')})"
+    letter = match.group(1).upper()
+    media_path = str(Path(media_root) / row["path"].removeprefix("./"))
+    data_type = row["data_type"]
+    assert data_type in _PLACEHOLDER, f"unsupported data_type: {data_type}"
+    return {
+        "prompt": _PROMPT_TEMPLATE.format(
+            placeholder=_PLACEHOLDER[data_type],
+            problem=row["problem"].strip(),
+            options="\n".join(row["options"]),
+        ),
+        "videos": [media_path] if data_type == "video" else [],
+        "images": [media_path] if data_type == "image" else [],
+        "label": letter,
+        "metadata": {"choices": list(row["options"])},
+    }
+
+
+def convert_file(
+    src,
+    dst,
+    media_root,
+    data_types: tuple[str, ...] = ("video", "image"),
+    path_prefixes: tuple[str, ...] | None = None,
+    problem_types: tuple[str, ...] = ("multiple choice",),
+    require_media: bool = False,
+    max_samples: int | None = None,
+) -> int:
+    with open(src) as f:
+        rows = json.load(f)
+
+    n = 0
+    with open(dst, "w") as fout:
+        for row in rows:
+            if max_samples is not None and n >= max_samples:
+                break
+            if row["data_type"] not in data_types or row["problem_type"] not in problem_types:
+                continue
+            rel = row["path"].removeprefix("./")
+            if path_prefixes and not rel.startswith(tuple(path_prefixes)):
+                continue
+            if require_media and not (Path(media_root) / rel).exists():
+                continue
+            fout.write(json.dumps(convert_row(row, media_root)) + "\n")
+            n += 1
+    assert n > 0, f"no rows converted from {src}"
+    return n
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--src", required=True, help="Video-R1-260k.json")
+    parser.add_argument("--dst", required=True)
+    parser.add_argument("--media-root", required=True, help="dir the media zips were extracted into")
+    parser.add_argument("--data-types", nargs="+", default=["video", "image"])
+    parser.add_argument("--path-prefixes", nargs="+", default=None, help="e.g. CLEVRER to match downloaded zips")
+    parser.add_argument("--require-media", action="store_true", help="skip rows whose media file is absent")
+    parser.add_argument("--max-samples", type=int, default=None)
+    args = parser.parse_args()
+    n = convert_file(
+        args.src,
+        args.dst,
+        media_root=args.media_root,
+        data_types=tuple(args.data_types),
+        path_prefixes=tuple(args.path_prefixes) if args.path_prefixes else None,
+        require_media=args.require_media,
+        max_samples=args.max_samples,
+    )
+    print(f"[done] {n} samples -> {args.dst}")
+
+
+if __name__ == "__main__":
+    main()