class_docstr: Dataset for RETRO model.
  
  It constructs single data record from the training/retrieval indexed retrieval dataset and knn index file.
  The KNN index file maps data chunk id to K-nearest neighbors in the the retrieval dataset chunk ids.
  First, it loads a long sequence (2048) from training dataset. Then for each chunk in the sequence, it finds the kNN 
  chunks from the retrieval dataset using the KNN index. Lastly, compute the masks based on pad id.
class_docstr_tok: ["Dataset","for","RETRO","model",".","It","constructs","single","data","record","from","the","training\/retrieval","indexed","retrieval","dataset","and","knn","index","file",".","The","KNN","index","file","maps","data","chunk","id","to","K-nearest","neighbors","in","the","the","retrieval","dataset","chunk","ids",".","First",",","it","loads","a","long","sequence","(","2048",")","from","training","dataset",".","Then","for","each","chunk","in","the","sequence",",","it","finds","the","kNN","chunks","from","the","retrieval","dataset","using","the","KNN","index",".","Lastly",",","compute","the","masks","based","on","pad","id","."]
code: def __init__(
      self,
      cfg,
      trainer,
      tokenizer,
      name: str,
      data_prefix: str,
      documents,  # document ids in the indexed_dataset used for this dataset
      indexed_dataset: MMapRetrievalIndexedDataset,
      num_samples: int,  # number of data samples,  max_steps * global_batch_size
      seq_length: int,  # input seq length
      seed: int,
      knn_index: KNNIndex,
      retrieval_index: MMapRetrievalIndexedDataset,
  ):
      if not HAVE_MEGATRON_CORE:
          raise ImportError(
              "megatron-core was not found. Please see the NeMo README for installation instructions: https://github.com/NVIDIA/NeMo#megatron-gpt."
          )
  
      super().__init__()
      self.name = name
      self.indexed_dataset: MMapRetrievalIndexedDataset = (
          indexed_dataset
      )
      self.knn_index: KNNIndex = knn_index
      self.retrieval_index: MMapRetrievalIndexedDataset = (
          retrieval_index
      )
      self.chunk_size = self.indexed_dataset.chunk_size
  
      # make sure seq_length is a multiple of chunk_size
      assert seq_length % self.chunk_size == 0
      # Checks
      assert np.min(documents) >= 0
      assert np.max(documents) < indexed_dataset.sizes.shape[0]
  
      self.eos_id = tokenizer.eos_id
      self.pad_id = tokenizer.pad_id
  
      assert self.retrieval_index._index.retrieval_db
      self._validate_pad_id()
  
      # save index mappings to a configurable dir
      self.index_mapping_dir = cfg.data.get("index_mapping_dir", None)
      self.neighbors = cfg.data.get("neighbors", self.knn_index.K)
      # the number of neighbors cannot exceed the max number of neighbors in the index
      assert self.neighbors <= self.knn_index.K
      # create index_mapping_dir on rank 0
      if (
          torch.distributed.is_available()
          and torch.distributed.is_initialized()
      ):
          if torch.distributed.get_rank() == 0:
              if (
                  self.index_mapping_dir is not None
                  and not os.path.isdir(self.index_mapping_dir)
              ):
                  os.makedirs(self.index_mapping_dir)
          torch.distributed.barrier()
  
      # Build index mappings.
      (
          self.doc_idx,
          self.sample_idx,
          self.shuffle_idx,
      ) = _build_index_mappings(
          self.name,
          data_prefix,
          documents,
          self.indexed_dataset.sizes,
          num_samples,
          seq_length,
          seed,
          index_mapping_dir=self.index_mapping_dir,
      )
      if len(self.doc_idx) > np.iinfo("int32").max:
          raise "number of epochs exceeds the maximum number for int32 used by sample_idx"
      self.padding_context = (
          np.ones(
              2 * self.chunk_size,
              dtype=self.retrieval_index._index.dtype,
          )
          * self.pad_id
      )
  
imports: import os
  from typing import List
  import numpy
  import torch
  from nemo.collections.nlp.data.language_modeling.megatron.base_dataset_utils import get_datasets_weights_and_num_samples, get_train_valid_test_split_
  from nemo.collections.nlp.data.language_modeling.megatron.blendable_dataset import BlendableDataset
  from nemo.collections.nlp.data.language_modeling.megatron.gpt_dataset import _build_index_mappings, get_indexed_dataset_
  from nemo.collections.nlp.data.language_modeling.megatron.indexed_retrieval_dataset import KNNIndex, MMapRetrievalIndexedDataset
  from nemo.core import Dataset
  from nemo.utils import logging
label_desc: Use image node_id 1 to create a new RETRODataset object from inherited base classes: Dataset with example: obj = RETRODataset(cfg, trainer, tokenizer, name, data_prefix, documents, indexed_dataset, num_samples, seq_length, seed, knn_index, retrieval_index)