Change the text name of sample for compatible with Huggingface trainer (bigscience-workshop#289)

Signed-off-by: yuanwu <[email protected]>

Author: yuanwu2017

megatron/data/gpt_dataset.py

@@ -301,14 +301,21 @@ def __getitem__(self, idx):
                 self.doc_idx[doc_index_l],
                 length=offset_l + 1))
             sample = np.concatenate(sample_list)
+
+        text_name = 'text'
+        if args.use_dataset_only:
+            text_name = 'input_ids'
+        sample_dict = {text_name: np.array(sample, dtype=np.int64)}
         if args.return_data_index:
-            return {'text': np.array(sample, dtype=np.int64),
-                    'index': np.array([orig_idx], dtype=np.int64)}
-        elif self.return_doc_ids: # for retro preprocessing
-            return {'text': np.array(sample, dtype=np.int64),
-                    'doc_ids': np.array(doc_ids, dtype=np.int64)}
-        else:
-            return {'text': np.array(sample, dtype=np.int64)}
+            sample_dict.update({'index': np.array([orig_idx], dtype=np.int64)})
+
+        if self.return_doc_ids: # for retro preprocessing
+            sample_dict.update({'doc_ids': np.array(doc_ids, dtype=np.int64)})
+
+        if args.use_dataset_only:
+            sample_dict.update({'labels': np.array(sample, dtype=np.int64)})
+
+        return sample_dict
 
 
 def _build_index_mappings(name, data_prefix, documents, sizes,