code: def test_knn_index(self):
      data_file = "/tmp/test"
      index_file = data_file + ".idx"
      K = 8
      index_files = [f"{data_file}_{i}.idx" for i in range(3)]
      merged_file = "/tmp/merged.idx"
      try:
          with KNNIndex.writer(index_file, K) as w:
              map_np0 = np.random.randint(0, 100, (50, K))
              w.write(map_np0)
              map_np1 = np.random.randint(0, 100, (50, K))
              w.write(map_np1)
              map_np2 = np.random.randint(0, 100, (50, K))
              w.write(map_np2)
          f = KNNIndex(index_file)
          assert f.K == K
          assert (
              f.len
              == map_np0.shape[0] + map_np1.shape[0] + map_np2.shape[0]
          )
          assert np.array_equal(map_np0, f.knn_map[:50])
          assert np.array_equal(map_np1, f.knn_map[50:100])
          assert np.array_equal(map_np2, f.knn_map[100:])
          assert np.array_equal(f.get_KNN_chunk_ids(5), map_np0[5])
          assert f.chunk_start_id == 0
          assert f.chunk_end_id == f.len
  
          with KNNIndex.writer(index_file, K, 100) as w:
              map_np0 = np.random.randint(0, 100, (50, K))
              w.write(map_np0)
              map_np1 = np.random.randint(0, 100, (50, K))
              w.write(map_np1)
              map_np2 = np.random.randint(0, 100, (50, K))
              w.write(map_np2)
          f = KNNIndex(index_file)
          assert f.K == K
          assert (
              f.len
              == map_np0.shape[0] + map_np1.shape[0] + map_np2.shape[0]
          )
          assert np.array_equal(map_np0, f.knn_map[:50])
          assert np.array_equal(map_np1, f.knn_map[50:100])
          assert np.array_equal(map_np2, f.knn_map[100:])
          assert np.array_equal(
              f.get_KNN_chunk_ids(5 + 100), map_np0[5]
          )
          assert f.chunk_start_id == 100
          assert f.chunk_end_id == f.len + 100
  
          # test multiple sharding indices
          inputs = []
          start = 0
          for i in range(3):
              with KNNIndex.writer(
                  index_files[i], K, offset=start
              ) as w:
                  map_np0 = np.random.randint(0, 100, (50, K))
                  inputs.append(map_np0)
                  w.write(map_np0)
                  map_np1 = np.random.randint(0, 100, (50, K))
                  inputs.append(map_np1)
                  w.write(map_np1)
              f = KNNIndex(index_files[i])
              start += f.len
          merge_knn_files(index_files, merged_file)
          f = KNNIndex(merged_file)
          input_array = np.vstack(inputs)
          assert f.len == 100 * 3
          for i in range(300):
              assert np.array_equal(
                  f.get_KNN_chunk_ids(i), input_array[i]
              )
          assert f.chunk_start_id == 0
          assert f.chunk_end_id == f.len
          assert f.K == K
  
      finally:
          os.remove(index_file)
          for i in range(3):
              os.remove(index_files[i])
          os.remove(merged_file)
  
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 7 for calling the TestRetrievalIndexFiles obj's underlying member method code with example usage: obj.test_knn_index() without return types