Fix inconsistent merge (#30)

* I suspect this was wrong

* Take in account offset

* Woops

Co-authored-by: Thomas <ö95242+thomasw21@users.noreply.github.com>

megatron/data/indexed_dataset.py

@@ -298,13 +298,16 @@ def merge_file_(self, another_file):
         index = IndexedDataset(another_file)
         assert index.dtype == self.dtype
 
+        offset = len(self.sizes)
+
         begin = self.data_offsets[-1]
         for offset in index.data_offsets[1:]:
             self.data_offsets.append(begin + offset)
         self.sizes.extend(index.sizes)
         begin = self.dim_offsets[-1]
         for dim_offset in index.dim_offsets[1:]:
             self.dim_offsets.append(begin + dim_offset)
+        self.doc_idx.extend( (offset + index.doc_idx)[1:] )
 
         with open(data_file_path(another_file), 'rb') as f:
             while True:
@@ -566,8 +569,9 @@ def merge_file_(self, another_file):
         total_len = len(index.sizes)+len(self._sizes)
         print(f"    concat {another_file} size={len(index.sizes)} for a total size of {total_len}")
 
-        for size in index.sizes:
-            self._sizes.append(size)
+        offset = len(self._sizes)
+        self._sizes.extend(index.sizes)
+        self._doc_idx.extend( (offset + index.doc_idx)[1:] )
 
         # Concatenate data
         with open(data_file_path(another_file), 'rb') as f: