code: def test_retro_dataset_stride32(self):
      chunk_size = 64
      pad_id = 0
      sentence1 = torch.arange(0, 200, 2, dtype=torch.int64)
      sentence2 = torch.arange(1, 500, 2, dtype=torch.int64)
      sentence3 = torch.arange(0, 300, 2, dtype=torch.int64)
      sentence4 = torch.arange(1, 400, 2, dtype=torch.int64)
  
      # test the case that
      # training data and retrieval data are different
  
      data_file = "/tmp/test_data"
      data_index_file = data_file + ".idx"
      data_bin_file = data_file + ".bin"
      db_file = "/tmp/test_db_data"
      db_index_file = db_file + ".idx"
      db_bin_file = db_file + ".bin"
      K = 8
      map_index_file = "/tmp/test_map.idx"
      index_path = "/tmp"
  
      cfg = OmegaConf.create(
          {"data": {"index_mapping_dir": index_path}}
      )
  
      # dummy tokenizer
      class Tokenizer:
          eos_id = 1
          pad_id = 0
  
      tokenizer = Tokenizer()
  
      num_samples = 100
      stride = 32
      seq_len = 192
      name = "test"
      data_prefix = "pref"
      seed = 1
      _filename = index_path + "/" + data_prefix
      _filename += "_{}_indexmap".format(name)
      _filename += "_{}ns".format(num_samples)
      _filename += "_{}sl".format(seq_len)
      _filename += "_{}s".format(seed)
      doc_idx_filename = _filename + "_doc_idx.npy"
      sample_idx_filename = _filename + "_sample_idx.npy"
      shuffle_idx_filename = _filename + "_shuffle_idx.npy"
  
      try:
          builder = MMapRetrievalIndexedDatasetBuilder(
              data_bin_file, chunk_size, pad_id, False, stride=32
          )
          builder.add_item(sentence1)
          builder.add_item(sentence2)
          builder.finalize(data_index_file)
  
          builder = MMapRetrievalIndexedDatasetBuilder(
              db_bin_file, chunk_size, pad_id, True, stride=32
          )
          builder.add_item(sentence3)
          builder.add_item(sentence4)
          builder.finalize(db_index_file)
  
          # load the data
          data_index = MMapRetrievalIndexedDataset(data_file)
          db_index = MMapRetrievalIndexedDataset(db_file)
  
          with KNNIndex.writer(map_index_file, K) as w:
              map_np = np.random.randint(
                  -3, db_index.chunks, (data_index.chunks, K)
              )
              w.write(map_np)
          map_index = KNNIndex(map_index_file)
  
          documents = np.arange(0, data_index.sizes.shape[0])
          d = RETRODataset(
              cfg,
              None,
              tokenizer,
              name,
              data_prefix,
              documents,
              data_index,
              num_samples,
              seq_len,
              seed,
              map_index,
              db_index,
          )
          for i in range(len(d)):
              record = d[i]
              assert record["tokens"].shape[0] == seq_len
              assert record["labels"].shape[0] == seq_len
              assert (
                  record["retrieved_ids"].shape[0]
                  == seq_len // chunk_size
              )
              assert record["retrieved_ids"].shape[1] == K
              assert record["retrieved_ids"].shape[2] == chunk_size * 2
              assert record["tokens_mask"].shape[0] == seq_len
  
      finally:
          os.remove(data_bin_file)
          os.remove(data_index_file)
          os.remove(db_bin_file)
          os.remove(db_index_file)
          os.remove(map_index_file)
          os.remove(doc_idx_filename)
          os.remove(sample_idx_filename)
          os.remove(shuffle_idx_filename)
  
      # test the case that
      # training data and retrieval data are the same
  
      try:
          builder = MMapRetrievalIndexedDatasetBuilder(
              db_bin_file, chunk_size, pad_id, True, stride=32
          )
          builder.add_item(sentence1)
          builder.add_item(sentence2)
          builder.add_item(sentence3)
          builder.add_item(sentence4)
          builder.finalize(db_index_file)
  
          # load the data
          data_index = MMapRetrievalIndexedDataset(db_file)
          db_index = MMapRetrievalIndexedDataset(db_file)
  
          with KNNIndex.writer(map_index_file, K) as w:
              map_np = np.random.randint(
                  -3, db_index.chunks, (data_index.chunks, K)
              )
              w.write(map_np)
          map_index = KNNIndex(map_index_file)
  
          documents = np.arange(0, data_index.sizes.shape[0])
          d = RETRODataset(
              cfg,
              None,
              tokenizer,
              name,
              data_prefix,
              documents,
              data_index,
              num_samples,
              seq_len,
              seed,
              map_index,
              db_index,
          )
          for i in range(len(d)):
              record = d[i]
              assert record["tokens"].shape[0] == seq_len
              assert record["labels"].shape[0] == seq_len
              assert (
                  record["retrieved_ids"].shape[0]
                  == seq_len // chunk_size
              )
              assert record["retrieved_ids"].shape[1] == K
              assert record["retrieved_ids"].shape[2] == chunk_size * 2
              assert record["tokens_mask"].shape[0] == seq_len
  
      finally:
          os.remove(db_bin_file)
          os.remove(db_index_file)
          os.remove(map_index_file)
          os.remove(doc_idx_filename)
          os.remove(sample_idx_filename)
          os.remove(shuffle_idx_filename)
  
imports: import os
  import numpy
  import pytest
  import torch
  from numpy.testing import assert_array_equal
  from omegaconf import OmegaConf
  from scripts.nlp_language_modeling.build_knn_map_index import build_map, dedup
  from nemo.collections.nlp.data.language_modeling.megatron.indexed_retrieval_dataset import KNNIndex, MMapRetrievalIndexedDataset, MMapRetrievalIndexedDatasetBuilder, merge_knn_files
  from nemo.collections.nlp.data.language_modeling.megatron.retro_dataset import RETRODataset
label_desc: Use image node_id 9 for calling the TestRetrievalIndexFiles obj's underlying member method code with example usage: obj.test_retro_dataset_stride32() without return types