diff --git a/internvl_chat_gpt_oss/internvl/train/internvl_chat_finetune.py b/internvl_chat_gpt_oss/internvl/train/internvl_chat_finetune.py index 01145fdc..1d8e6864 100644 --- a/internvl_chat_gpt_oss/internvl/train/internvl_chat_finetune.py +++ b/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 diff --git a/tests/test_gpt_oss_packed_preprocessing.py b/tests/test_gpt_oss_packed_preprocessing.py new file mode 100644 index 00000000..f60c4bdb --- /dev/null +++ b/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()