[data] feat: Add dataset for Qwen-Image (verl-project#6)

* add entroypoint (verl-project#1)

* add training engine (verl-project#2)

* add training engine

* fix init

* fix typs

* move folders & make for two-forward pass in training loop (verl-project#4)

* Add diffusion reward loop (verl-project#3)

* init reward; add ocr reward

* update disrm input

* add unit test

* pass ut

* fix typos/bugs

* update copyright

* [fix] update customized reward func in UT (verl-project#5)

* init reward; add ocr reward

* update disrm input

* add unit test

* pass ut

* fix typos/bugs

* update copyright

* update customized reward_fn

* init dataset for Qwen-Image

* pass UT

* update return, update UT

* pass UT

* align with rl_dataset

* pass UT

* update filter long prompts

* debug

* clean code

---------

Co-authored-by: Cheung Ka Wai <zhtmike@gmail.com>

Author: chenyingshu

tests/experimental/reward_loop/test_diffusion_reward_model_genrm.py

@@ -47,7 +47,7 @@ def create_data_samples(tokenizer) -> DataProto:
 
     data = DataProto.from_dict(
         tensors={
-            "prompts": prompt_ids,
+            "input_ids": prompt_ids,
             "responses": responses,
         },
         non_tensors={

tests/utils/dataset/test_qwen_dataset_on_cpu.py

@@ -0,0 +1,80 @@
+# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import os
+
+import torch
+from omegaconf import OmegaConf
+from torch.utils.data import DataLoader
+
+from verl import DataProto
+from verl.utils import hf_tokenizer
+from verl.utils.dataset import QwenDataset
+from verl.utils.dataset.rl_dataset import collate_fn
+
+
+def get_ocr_data():
+    # prepare test dataset
+    local_folder = os.path.expanduser("~/data/ocr/")
+    local_path = os.path.join(local_folder, "train.txt")
+    os.makedirs(local_folder, exist_ok=True)
+    return local_path
+
+
+def test_qwen_dataset():
+    tokenizer = hf_tokenizer(os.path.expanduser("~/models/Qwen/Qwen-Image"), trust_remote_code=True)
+    local_path = get_ocr_data()
+    config = OmegaConf.create(
+        {
+            "max_prompt_length": 1024,
+            "filter_overlong_prompts": True,
+            "data_source": "ocr",
+        }
+    )
+    dataset = QwenDataset(data_files=local_path, tokenizer=tokenizer, config=config)
+
+    dataloader = DataLoader(dataset=dataset, batch_size=16, shuffle=True, drop_last=True, collate_fn=collate_fn)
+
+    a = next(iter(dataloader))
+
+    tensors = {}
+    non_tensors = {}
+
+    for key, val in a.items():
+        if isinstance(val, torch.Tensor):
+            tensors[key] = val
+        else:
+            non_tensors[key] = val
+
+    data_proto = DataProto.from_dict(tensors=tensors, non_tensors=non_tensors)
+    assert len(data_proto) == 16
+    assert "input_ids" in data_proto.batch
+    assert "attention_mask" in data_proto.batch
+
+
+def test_qwen_dataset_with_max_samples():
+    tokenizer = hf_tokenizer(os.path.expanduser("~/models/Qwen/Qwen-Image"), trust_remote_code=True)
+    local_path = get_ocr_data()
+    config = OmegaConf.create(
+        {
+            "max_prompt_length": 1024,
+            "filter_overlong_prompts": True,
+            "data_source": "ocr",
+        }
+    )
+    dataset = QwenDataset(data_files=local_path, tokenizer=tokenizer, config=config, max_samples=5)
+    assert len(dataset) == 5
+
+    # test split
+    dataset_split = dataset.split(5)
+    assert len(dataset_split) == 5

verl/experimental/reward_loop/reward_manager/diffusion.py

@@ -45,7 +45,6 @@ def __init__(self, config, tokenizer, compute_score=None, reward_router_address=
     async def run_single(self, data: DataProto) -> dict:
         assert len(data) == 1, "Only support single data item"
         data_item = data[0]
-        # prompt_str = self.tokenizer.decode(data_item.batch["prompts"], skip_special_tokens=True)
         response_image = data_item.batch["responses"]
         data_source = data_item.non_tensor_batch["data_source"]
         ground_truth = data_item.non_tensor_batch["reward_model"]["ground_truth"]

verl/utils/dataset/__init__.py

@@ -12,8 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from .qwen_dataset import QwenDataset
 from .rl_dataset import RLHFDataset
 from .rm_dataset import RMDataset
 from .sft_dataset import SFTDataset
 
-__all__ = ["RLHFDataset", "RMDataset", "SFTDataset"]
+__all__ = ["RLHFDataset", "RMDataset", "SFTDataset", "QwenDataset"]