feature(nyz): add vision input support and fix bugs

Author: PaParaZz1

ding/utils/data/rlhf_offline_dataset.py

@@ -30,6 +30,7 @@ def __init__(
             tokenizer,
             max_length: int,
             input_key: str = "input",
+            extra_input_keys: List[str] = [],
             output_key: str = "output",
             label_key: str = "label",
             apply_chat_template: bool = False,
@@ -41,6 +42,7 @@ def __init__(
         super().__init__()
         self.tokenizer = tokenizer
         self.max_length = max_length
+        self.extra_input_keys = extra_input_keys
 
         if apply_chat_template:
             apply_chat_template = self.tokenizer.apply_chat_template
@@ -53,6 +55,7 @@ def __init__(
                 self._preprocess_data,
                 input_template=input_template,
                 input_key=input_key,
+                extra_input_keys=extra_input_keys,
                 output_key=output_key,
                 label_key=label_key,
                 apply_chat_template=apply_chat_template
@@ -67,29 +70,40 @@ def __init__(
             self.responses = processed_dataset["response"]
             self.labels = processed_dataset["label"]
             self.prompt_ids_lens = processed_dataset["prompt_ids_len"]
+            for key in extra_input_keys:
+                setattr(self, key, processed_dataset[key])
         else:
             self.prompts = []
             self.responses = []
             self.labels = []
             self.prompt_ids_lens = []
+            for key in extra_input_keys:
+                setattr(self, key, [])
             for data in tqdm(dataset, desc="Preprocessing data", disable=not get_rank() == 0):
                 processed_data = self._preprocess_data(data)
                 if processed_data["prompt"] is not None:
                     self.prompts.append(processed_data["prompt"])
                     self.responses.append(processed_data["response"])
                     self.labels.append(processed_data["label"])
                     self.prompt_ids_lens.append(processed_data["prompt_ids_len"])
+                    for key in extra_input_keys:
+                        getattr(self, key).append(processed_data[key])
 
     def _preprocess_data(
             self,
             data: Dict[str, Any],
             input_template: str = None,
             input_key: str = "input",
+            extra_input_keys: List[str] = [],
             output_key: str = "output",
             label_key: str = "label",
             apply_chat_template: Union[bool, Callable] = False,
     ) -> str:
         label = data[label_key]
+        if extra_input_keys:
+            extra_inputs = {key: data[key] for key in extra_input_keys}
+        else:
+            extra_inputs = {}
 
         if apply_chat_template:
             if output_key:
@@ -120,7 +134,13 @@ def _preprocess_data(
         if prompt_ids_len >= self.max_length - 2:
             prompt = None
 
-        return {"prompt": prompt, "response": response, "label": label, "prompt_ids_len": prompt_ids_len}
+        return {
+            "prompt": prompt,
+            "response": response,
+            "label": label,
+            "prompt_ids_len": prompt_ids_len,
+            **extra_inputs
+        }
 
     def __len__(self) -> int:
         """
@@ -135,14 +155,21 @@ def __getitem__(self, idx: int) -> Dict[str, Union[torch.Tensor, int]]:
         """
         Overview:
             Get the item at the given index.
+        Arguments:
+            - idx (int): The index of the item to get.
         Returns:
             - item (Dict[str, Union[torch.Tensor, int]]): The item at the given index.
         """
+        if self.extra_input_keys:
+            extra_inputs = {key: getattr(self, key)[idx] for key in self.extra_input_keys}
+        else:
+            extra_inputs = {}
         return {
             "prompt": self.prompts[idx],
             "response": self.responses[idx],
             "label": self.labels[idx],
-            "prompt_ids_len": self.prompt_ids_lens[idx]
+            "prompt_ids_len": self.prompt_ids_lens[idx],
+            **extra_inputs
         }
 
     def collate_fn(self, item_list: List[Dict[str, Union[torch.Tensor, int]]]):
@@ -164,13 +191,17 @@ def tokenizer(prompt: str, response: str):
             inputs["attention_mask"][0][-1] = True
             return inputs["input_ids"], inputs["attention_mask"]
 
-        tot_ids, tot_masks, tot_labels, prompt_ids_lens = [], [], [], []
+        tot_ids, tot_masks, tot_labels, prompt_ids_lens, tot_extra_inputs = [], [], [], [], {}
         for item in item_list:
             input_ids, attention_mask = tokenizer(item["prompt"], item["response"])
             tot_ids.append(input_ids)
             tot_masks.append(attention_mask)
             tot_labels.append(item["label"])
             prompt_ids_lens.append(item["prompt_ids_len"])
+            for key in self.extra_input_keys:
+                if key not in tot_extra_inputs:
+                    tot_extra_inputs[key] = []
+                tot_extra_inputs[key].append(item[key])
 
         # add unmatched y'| x (used to estimate the KL divergence between policy and reference)
         for idx in range(len(item_list)):
@@ -180,7 +211,11 @@ def tokenizer(prompt: str, response: str):
             tot_masks.append(attention_mask)
             tot_labels.append(-1)
             prompt_ids_lens.append(item_list[idx]["prompt_ids_len"])
+            for key in self.extra_input_keys:
+                if key not in tot_extra_inputs:
+                    tot_extra_inputs[key] = []
+                tot_extra_inputs[key].append(item_list[idx][key])
 
         input_ids = zero_pad_sequences(tot_ids, side="right", value=self.tokenizer.pad_token_id)
         attention_mask = zero_pad_sequences(tot_masks, side="right")
-        return input_ids, attention_mask, torch.LongTensor(tot_labels), prompt_ids_lens
+        return input_ids, attention_mask, torch.LongTensor(tot_labels), prompt_ids_lens, tot_extra_inputs

ding/utils/data/tests/test_rlhf_offline_dataset.py

@@ -3,6 +3,11 @@
 from ding.utils.data import OfflineRLDataset
 from transformers import AutoTokenizer
 
+IMG_CONTEXT_TOKEN = '<IMG_CONTEXT>'
+IMG_START_TOKEN = '<img>'
+IMG_END_TOKEN = '</img>'
+IMG_CONTEXT_NUM = 10  # user-defined number of image patches in the context
+
 
 @pytest.fixture
 def dataset():
@@ -11,16 +16,18 @@ def dataset():
     # split pair data into two separate datasets
     hf_dataset_1 = hf_dataset.map(
         lambda x: {
-            "prompt": x["query"],
+            "query": f"{IMG_START_TOKEN}{IMG_CONTEXT_TOKEN * IMG_CONTEXT_NUM}{IMG_END_TOKEN}\n{x['query']}",
+            "image": x["image"],
             "response": x["response"][0],
-            'human_ranking': x["human_ranking"][0]
+            "human_ranking": x["human_ranking"][0]
         }
     )
     hf_dataset_2 = hf_dataset.map(
         lambda x: {
-            "prompt": x["query"],
+            "query": f"{IMG_START_TOKEN}{IMG_CONTEXT_TOKEN * IMG_CONTEXT_NUM}{IMG_END_TOKEN}\n{x['query']}",
+            "image": x["image"],
             "response": x["response"][1],
-            'human_ranking': x["human_ranking"][0]
+            "human_ranking": x["human_ranking"][1]
         }
     )
     # combine two datasets
@@ -33,7 +40,7 @@ def dataset():
 @pytest.fixture
 def tokenizer():
     # Load a tokenizer
-    return AutoTokenizer.from_pretrained("Qwen/Qwen2.5-Coder-7B")
+    return AutoTokenizer.from_pretrained("OpenGVLab/InternVL2_5-4B")
 
 
 @pytest.mark.unittest
@@ -44,6 +51,7 @@ def test_offline_rl_dataset_initialization(dataset, tokenizer):
         tokenizer=tokenizer,
         max_length=1024,
         input_key="query",
+        extra_input_keys=["image"],
         output_key="response",
         label_key="human_ranking"
     )
@@ -53,6 +61,7 @@ def test_offline_rl_dataset_initialization(dataset, tokenizer):
         tokenizer=tokenizer,
         max_length=256,
         input_key="query",
+        extra_input_keys=["image"],
         output_key="response",
         label_key="human_ranking"
     )
@@ -68,6 +77,7 @@ def test_offline_rl_dataset_item_retrieval(dataset, tokenizer):
         tokenizer=tokenizer,
         max_length=256,
         input_key="query",
+        extra_input_keys=["image"],
         output_key="response",
         label_key="human_ranking"
     )
@@ -76,6 +86,7 @@ def test_offline_rl_dataset_item_retrieval(dataset, tokenizer):
     assert "response" in item
     assert "label" in item
     assert "prompt_ids_len" in item
+    assert "image" in item
     print(item)
 
 
@@ -92,8 +103,11 @@ def test_offline_rl_dataset_collate_fn(dataset, tokenizer):
     )
     B = 10
     item_list = [offline_dataset[i] for i in range(B)]
-    input_ids, attention_mask, labels, prompt_ids_lens = offline_dataset.collate_fn(item_list)
+    input_ids, attention_mask, labels, prompt_ids_lens, extra_inputs = offline_dataset.collate_fn(item_list)
     assert input_ids.size(0) == len(item_list) * 2  # because of the unmatched y'| x
     assert attention_mask.size(0) == len(item_list) * 2
     assert labels.size(0) == len(item_list) * 2
     assert len(prompt_ids_lens) == len(item_list) * 2
+    for key in offline_dataset.extra_input_keys:
+        assert key in extra_inputs
+        assert extra_inputs[key].size(0) == len(item_list) * 2