class_docstr: Index file for fast KNN mapping.
  It is built by `build_knn_map_index.py` script.
  It contains a big matrix of shape (chunk_id, K neighbors)
  where `chunk_id` are all the chunk ids in the RETRO training data.
  E.g. the KNN neighbor chunk ids in the retrieval data for ith chunk id in the training data
  is self.knn_map[i].
  This index can hold partial maps used for building sharding index.
class_docstr_tok: ["Index","file","for","fast","KNN","mapping",".","It","is","built","by","`","build_knn_map_index.py","`","script",".","It","contains","a","big","matrix","of","shape","(","chunk_id",",","K","neighbors",")","where","`","chunk_id","`","are","all","the","chunk","ids","in","the","RETRO","training","data",".","E.g",".","the","KNN","neighbor","chunk","ids","in","the","retrieval","data","for","ith","chunk","id","in","the","training","data","is","self.knn_map","[","i","]",".","This","index","can","hold","partial","maps","used","for","building","sharding","index","."]
code: def __init__(self, path, skip_warmup=True):
      with open(path, "rb") as stream:
          magic_test = stream.read(9)
          assert (
              self._HDR_MAGIC == magic_test
          ), "Index file doesn't match expected format. "
          version = struct.unpack("<Q", stream.read(8))
          assert (1,) == version
  
          self.K = struct.unpack("<Q", stream.read(8))[0]
          self.len = struct.unpack("<Q", stream.read(8))[0]
          self.chunk_start_id = struct.unpack("<Q", stream.read(8))[0]
          self.chunk_end_id = self.chunk_start_id + self.len
          offset = stream.tell()
  
      if not skip_warmup:
          logging.info("    warming up index mmap file...")
          _warmup_mmap_file(path)
  
      self._bin_buffer_mmap = np.memmap(path, mode="r", order="C")
      self._bin_buffer = memoryview(self._bin_buffer_mmap)
      logging.info("    reading KNN map")
      self.knn_map = np.frombuffer(
          self._bin_buffer,
          dtype=np.int64,
          count=self.len * self.K,
          offset=offset,
      ).reshape(self.len, self.K)
  
imports: import os
  import shutil
  import struct
  from functools import lru_cache
  from itertools import accumulate
  from typing import List
  import numpy
  import torch
  from nemo.utils import logging
label_desc: Use image node_id 2 to create a new KNNIndex object from inherited base classes: object with example: obj = KNNIndex(path, skip_warmup)