File size: 2,440 Bytes
53c10a4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 | from __future__ import annotations
import json
import pickle
import pytest
pytest.importorskip("torch")
pytest.importorskip("qwen_vl_utils")
from verl.utils.dataset import ( # noqa: E402
LocalJsonlDataset,
_align_media_placeholders,
)
def test_local_jsonl_preserves_task_specific_columns(tmp_path) -> None:
path = tmp_path / "mixed.jsonl"
rows = [
{
"problem": "temporal prompt",
"answer": "<answer>1 to 2</answer>",
"problem_type": "temporal grounding",
"videos": ["/media/a.mp4"],
"pred_span": [1.0, 2.0],
},
{
"problem": "segmentation prompt",
"answer": "<answer>{}</answer>",
"problem_type": "segmentation",
"videos": ["/media/b.mp4"],
"segmentation_output": {"object_id": 1},
},
]
path.write_text(
"".join(json.dumps(row) + "\n" for row in rows),
encoding="utf-8",
)
dataset = LocalJsonlDataset(str(path))
assert len(dataset) == 2
assert dataset[0]["pred_span"] == [1.0, 2.0]
assert dataset[1]["segmentation_output"] == {"object_id": 1}
assert "segmentation_output" not in dataset[0]
def test_local_jsonl_filter_and_pickle_keep_random_access(tmp_path) -> None:
path = tmp_path / "mixed.jsonl"
path.write_text(
"\n".join(
[
json.dumps({"problem_type": "tracking", "value": 1}),
"",
json.dumps({"problem_type": "video_qa_mc", "value": 2}),
]
)
+ "\n",
encoding="utf-8",
)
dataset = LocalJsonlDataset(str(path))
filtered = dataset.filter(lambda row: row["value"] == 2)
restored = pickle.loads(pickle.dumps(filtered))
assert len(dataset) == 2
assert len(restored) == 1
assert restored[0] == {"problem_type": "video_qa_mc", "value": 2}
def test_surplus_media_placeholders_are_removed() -> None:
prompt = "<image> <image> <image> Question with <image> literal tail"
aligned = _align_media_placeholders(prompt, "<image>", media_count=2)
assert aligned.count("<image>") == 2
assert "Question with" in aligned
assert "literal tail" in aligned
def test_missing_media_placeholders_are_rejected() -> None:
with pytest.raises(ValueError, match="1 <video> placeholder"):
_align_media_placeholders("<video> Question", "<video>", media_count=2)
|