Complete implementation of InverseClozeDataset with IndexedDataset

megatron/data/ict_dataset.py

 import random
+import os
+import time
 
 import numpy as np
+import torch
 from torch.utils.data import Dataset
 
 from megatron import get_tokenizer
-
+from megatron import print_rank_0
+from megatron import mpu
+from megatron.data import helpers
 
 class InverseClozeDataset(Dataset):
     """Dataset containing sentences and various 'blocks' for an inverse cloze task."""
@@ -14,17 +19,8 @@ class InverseClozeDataset(Dataset):
         self.name = name
         self.seed = seed
         self.max_seq_length = max_seq_length
-
         self.indexed_dataset = indexed_dataset
-
-        self.samples_mapping = get_samples_mapping(self.indexed_dataset,
-                                                    data_prefix,
-                                                    num_epochs,
-                                                    max_num_samples,
-                                                    self.max_seq_length,
-                                                    short_seq_prob,
-                                                    self.seed,
-                                                    self.name)
+        self.short_seq_prob = short_seq_prob
 
         tokenizer = get_tokenizer()
         self.vocab_id_list = list(tokenizer.inv_vocab.keys())
@@ -35,11 +31,11 @@ class InverseClozeDataset(Dataset):
         self.pad_id = tokenizer.pad
 
     def __len__(self):
-        return self.samples_mapping.shape[0]
+        return self.indexed_dataset.doc_idx.shape[0]
 
     def __getitem__(self, idx):
         # get rng state corresponding to index (allows deterministic random pair)
-        rng = random.Random(idx + 1000)
+        rng = random.Random(idx + self.seed)
         np_rng = np.random.RandomState(seed=[rng.randint(0, 2**32-1) for _ in range(16)])
 
         # get seq length. Save 2 tokens for beginning and end
@@ -64,29 +60,23 @@ class InverseClozeDataset(Dataset):
 
     def get_sentence_split_doc(self, idx):
         """fetch document at index idx and split into sentences"""
-        document = self.indexed_dataset[idx]
-        if isinstance(document, dict):
-            document = document['text']
-        lines = document.split('\n')
-        return [line for line in lines if line]
-
-    def sentence_tokenize(self, sent, sentence_num=0):
-        """tokenize sentence and get token types"""
-        tokens = self.tokenizer.EncodeAsIds(sent).tokenization
-        str_type = 'str' + str(sentence_num)
-        token_types = [self.tokenizer.get_type(str_type).Id]*len(tokens)
-        return tokens, token_types
-
-    def concat_and_pad_tokens(self, tokens, token_types):
+        doc_start = self.indexed_dataset.doc_idx[idx]
+        doc_end = self.indexed_dataset.doc_idx[idx + 1]
+
+        doc_sentences_array = self.indexed_dataset[doc_start:doc_end]
+        doc_sentences = [list(arr) for arr in doc_sentences_array]
+
+        return doc_sentences
+
+    def concat_and_pad_tokens(self, tokens):
         """concat with special tokens and pad sequence to self.max_seq_length"""
         tokens = [self.cls_id] + tokens + [self.sep_id]
-        token_types = [token_types[0]] + token_types + [token_types[0]]
-
         assert len(tokens) <= self.max_seq_length
-        num_pad = max(0, self.max_seq_length - len(tokens))
+
+        num_pad = self.max_seq_length - len(tokens)
         pad_mask = [0] * len(tokens) + [1] * num_pad
         tokens += [self.pad_id] * num_pad
-        token_types += [token_types[0]] * num_pad
+        token_types = [0] * self.max_seq_length
         return tokens, token_types, pad_mask
 
     def get_input_and_context(self, target_seq_length, rng, np_rng):
@@ -102,26 +92,22 @@ class InverseClozeDataset(Dataset):
                 if not doc:
                     doc = None
 
-            # set up and tokenize the entire selected document
             num_sentences = len(doc)
             padless_max_len = self.max_seq_length - 2
 
             # select a random sentence from the document as input
             # TODO: consider adding multiple input sentences.
             input_sentence_idx = rng.randint(0, num_sentences - 1)
-            tokens, token_types = self.sentence_tokenize(doc[input_sentence_idx], 0)
-            input_tokens, input_token_types = tokens[:target_seq_length], token_types[:target_seq_length]
+            input_tokens = doc[input_sentence_idx][:target_seq_length]
             if not len(input_tokens) > 0:
                 continue
 
-            context_tokens, context_token_types = [], []
+            context_tokens = []
             # 10% of the time, the input sentence is left in the context.
             # The other 90% of the time, keep it out.
             if rng.random() < 0.1:
                 context_tokens = input_tokens.copy()
-                context_token_types = input_token_types.copy()
 
-            # parameters for examining sentences to add to the context
             view_preceding = True
             view_radius = 1
             while len(context_tokens) < padless_max_len:
@@ -129,15 +115,13 @@ class InverseClozeDataset(Dataset):
                 if view_preceding:
                     examine_idx = input_sentence_idx - view_radius
                     if examine_idx >= 0:
-                        new_tokens, new_token_types = self.sentence_tokenize(doc[examine_idx], 0)
+                        new_tokens = doc[examine_idx]
                         context_tokens = new_tokens + context_tokens
-                        context_token_types = new_token_types + context_token_types
                 else:
                     examine_idx = input_sentence_idx + view_radius
                     if examine_idx < num_sentences:
-                        new_tokens, new_token_types = self.sentence_tokenize(doc[examine_idx], 0)
+                        new_tokens = doc[examine_idx]
                         context_tokens += new_tokens
-                        context_token_types += new_token_types
                     view_radius += 1
                 view_preceding = not view_preceding
                 if view_radius > num_sentences:
@@ -145,15 +129,12 @@ class InverseClozeDataset(Dataset):
 
             # assemble the tokens and token types of the context
             context_tokens = context_tokens[:padless_max_len]
-            context_token_types = context_token_types[:padless_max_len]
             if not len(context_tokens) > 0:
                 continue
 
             # concatenate 'CLS' and 'SEP' tokens and add extra token types
-            input_tokens, input_token_types, input_pad_mask = self.concat_and_pad_tokens(
-                input_tokens, input_token_types)
-            context_tokens, context_token_types, context_pad_mask = self.concat_and_pad_tokens(
-                context_tokens, context_token_types)
+            input_tokens, input_token_types, input_pad_mask = self.concat_and_pad_tokens(input_tokens)
+            context_tokens, context_token_types, context_pad_mask = self.concat_and_pad_tokens(context_tokens)
 
             return (input_tokens, input_token_types, input_pad_mask), \
                    (context_tokens, context_token_types, context_pad_mask)
@@ -161,82 +142,3 @@ class InverseClozeDataset(Dataset):
             raise RuntimeError("Could not get a valid data point from InverseClozeDataset")
 
 
-def get_samples_mapping(indexed_dataset,
-                         data_prefix,
-                         num_epochs,
-                         max_num_samples,
-                         max_seq_length,
-                         short_seq_prob,
-                         seed,
-                         name):
-    if not num_epochs:
-        if not max_num_samples:
-            raise ValueError("Need to specify either max_num_samples "
-                             "or num_epochs")
-        num_epochs = np.iinfo(np.int32).max - 1
-    if not max_num_samples:
-        max_num_samples = np.iinfo(np.int64).max - 1
-
-    # Filename of the index mapping
-    indexmap_filename = data_prefix
-    indexmap_filename += '_{}_indexmap'.format(name)
-    if num_epochs != (np.iinfo(np.int32).max - 1):
-        indexmap_filename += '_{}ep'.format(num_epochs)
-    if max_num_samples != (np.iinfo(np.int64).max - 1):
-        indexmap_filename += '_{}mns'.format(max_num_samples)
-    indexmap_filename += '_{}msl'.format(max_seq_length)
-    indexmap_filename += '_{:0.2f}ssp'.format(short_seq_prob)
-    indexmap_filename += '_{}s'.format(seed)
-    indexmap_filename += '.npy'
-
-    # Build the indexed mapping if not exist.
-    if torch.distributed.get_rank() == 0 and \
-       not os.path.isfile(indexmap_filename):
-        print(' > WARNING: could not find index map file {}, building '
-              'the indices on rank 0 ...'.format(indexmap_filename))
-
-        # Make sure the types match the helpers input types.
-        assert indexed_dataset.doc_idx.dtype == np.int64
-        assert indexed_dataset.sizes.dtype == np.int32
-
-        # Build samples mapping
-        verbose = torch.distributed.get_rank() == 0
-        start_time = time.time()
-        print_rank_0(' > building sapmles index mapping for {} ...'.format(
-            name))
-        samples_mapping = helpers.build_mapping(
-            indexed_dataset.doc_idx,
-            indexed_dataset.sizes,
-            num_epochs,
-            max_num_samples,
-            max_seq_length-3, # account for added tokens
-            short_seq_prob,
-            seed,
-            verbose)
-        print_rank_0(' > done building sapmles index maping')
-        np.save(indexmap_filename, samples_mapping, allow_pickle=True)
-        print_rank_0(' > saved the index mapping in {}'.format(
-            indexmap_filename))
-        # Make sure all the ranks have built the mapping
-        print_rank_0(' > elasped time to build and save samples mapping '
-                     '(seconds): {:4f}'.format(
-                         time.time() - start_time))
-    # This should be a barrier but nccl barrier assumes
-    # device_index=rank which is not the case for model
-    # parallel case
-    counts = torch.cuda.LongTensor([1])
-    torch.distributed.all_reduce(counts, group=mpu.get_data_parallel_group())
-    assert counts[0].item() == torch.distributed.get_world_size(
-        group=mpu.get_data_parallel_group())
-
-    # Load indexed dataset.
-    print_rank_0(' > loading indexed mapping from {}'.format(
-        indexmap_filename))
-    start_time = time.time()
-    samples_mapping = np.load(indexmap_filename, allow_pickle=True)
-    print_rank_0('    loaded indexed file in {:3.3f} seconds'.format(
-        time.time() - start_time))
-    print_rank_0('    total number of samples: {}'.format(
-        samples_mapping.shape[0]))
-
-    return samples_mapping