[Fix] Forward packed dataset mode to GPT-OSS preprocessors

internvl_chat_gpt_oss/internvl/train/internvl_chat_finetune.py

@@ -463,6 +463,7 @@ def multi_modal_get_item(self, data_item):
         ret = preprocess_function(self.template_name, [deepcopy(data_item['conversations'])],
                                   self.tokenizer, [self.num_image_token * num_patches],
                                   group_by_length=self.group_by_length,
+                                  use_packed_ds=self.use_packed_ds,
                                   ds_name=self.ds_name)
 
         # Calculate position_ids for packed dataset
@@ -519,6 +520,7 @@ def multi_modal_multi_image_get_item(self, data_item):
         num_image_tokens = [self.num_image_token * num_tile for num_tile in num_tiles]
         ret = preprocess_function(self.template_name, [deepcopy(data_item['conversations'])],
                                   self.tokenizer, num_image_tokens, group_by_length=self.group_by_length,
+                                  use_packed_ds=self.use_packed_ds,
                                   ds_name=self.ds_name, num_image=num_image)
 
         # Calculate position_ids for packed dataset
@@ -579,6 +581,7 @@ def video_get_item(self, data_item):
         num_image_tokens = [self.num_image_token] * num_patches
         ret = preprocess_function(self.template_name, [deepcopy(data_item['conversations'])],
                                   self.tokenizer, num_image_tokens, group_by_length=self.group_by_length,
+                                  use_packed_ds=self.use_packed_ds,
                                   ds_name=self.ds_name, num_image=num_patches)
 
         # Calculate position_ids for packed dataset
@@ -622,7 +625,8 @@ def pure_text_get_item(self, data_item):
         # Preprocess the conversations and generate the return dictionary
         ret = preprocess_function(self.template_name, [deepcopy(data_item['conversations'])],
                                   self.tokenizer, [self.num_image_token * num_patches], text_only=True,
-                                  group_by_length=self.group_by_length, ds_name=self.ds_name)
+                                  group_by_length=self.group_by_length, use_packed_ds=self.use_packed_ds,
+                                  ds_name=self.ds_name)
 
         # Calculate position_ids for packed dataset
         position_ids = ret['attention_mask'].long().cumsum(-1) - 1

tests/test_gpt_oss_packed_preprocessing.py

@@ -0,0 +1,63 @@
+import ast
+import unittest
+from pathlib import Path
+
+
+TRAINING_SCRIPT = (
+    Path(__file__).parents[1]
+    / 'internvl_chat_gpt_oss'
+    / 'internvl'
+    / 'train'
+    / 'internvl_chat_finetune.py'
+)
+
+
+class PackedPreprocessingTests(unittest.TestCase):
+    def test_all_supervised_modalities_forward_packed_dataset_mode(self):
+        tree = ast.parse(TRAINING_SCRIPT.read_text(encoding='utf-8'))
+        dataset_class = next(
+            node
+            for node in tree.body
+            if isinstance(node, ast.ClassDef) and node.name == 'LazySupervisedDataset'
+        )
+        methods = {
+            node.name: node
+            for node in dataset_class.body
+            if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))
+        }
+
+        for method_name in (
+            'multi_modal_get_item',
+            'multi_modal_multi_image_get_item',
+            'video_get_item',
+            'pure_text_get_item',
+        ):
+            with self.subTest(method=method_name):
+                preprocess_calls = [
+                    node
+                    for node in ast.walk(methods[method_name])
+                    if isinstance(node, ast.Call)
+                    and isinstance(node.func, ast.Name)
+                    and node.func.id == 'preprocess_function'
+                ]
+                self.assertEqual(len(preprocess_calls), 1)
+
+                packed_keyword = next(
+                    (keyword for keyword in preprocess_calls[0].keywords if keyword.arg == 'use_packed_ds'),
+                    None,
+                )
+                self.assertIsNotNone(packed_keyword)
+                self.assertEqual(
+                    ast.dump(packed_keyword.value),
+                    ast.dump(
+                        ast.Attribute(
+                            value=ast.Name(id='self', ctx=ast.Load()),
+                            attr='use_packed_ds',
+                            ctx=ast.Load(),
+                        )
+                    ),
+                )
+
+
+if __name__ == '__main__':
+    unittest.main()