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","."]
docstr: path: file path of the index
  K: number of neighbors for a chunk
  offset: start chunk_id for shard index
docstr_tok: ["path",":","file","path","of","the","index","K",":","number","of","neighbors","for","a","chunk","offset",":","start","chunk_id","for","shard","index"]
code: def writer(cls, path, K, offset=0):
      """
      path: file path of the index
      K: number of neighbors for a chunk
      offset: start chunk_id for shard index
      """
  
      class _Writer(object):
          def __enter__(self):
              self._file = open(path, "wb")
              self._file.write(cls._HDR_MAGIC)
              self._file.write(struct.pack("<Q", 1))
              self._file.write(struct.pack("<Q", K))
              # reserve the space for total number of chunks
              self._file.write(struct.pack("<Q", 0))
              # chunk start
              self._file.write(struct.pack("<Q", offset))
              self.K = K
              self.count_chunks = 0
              self.path = path
              return self
  
          def write(self, chunk_knn: np.array):
              assert chunk_knn.dtype == np.int64
              assert chunk_knn.shape[1] == self.K
              self._file.write(chunk_knn.tobytes(order="C"))
              self.count_chunks += chunk_knn.shape[0]
  
          def __exit__(self, exc_type, exc_val, exc_tb):
              self._file.close()
              # Update the chunk size, Since total number of chunks is determined in the end
              _bin_buffer_mmap = np.memmap(
                  self.path,
                  mode="r+",
                  order="C",
                  shape=(9 + 8 + 8 + 8),
              )
              buffer = memoryview(_bin_buffer_mmap)
              len_array = np.frombuffer(
                  buffer, dtype=np.int64, count=1, offset=9 + 8 + 8
              )
              len_array[0] = self.count_chunks
              _bin_buffer_mmap.flush()
              _bin_buffer_mmap._mmap.close()
  
      return _Writer()
  
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 1 for calling the KNNIndex obj's underlying member method code with example usage: obj.writer(cls, path, K, offset) and returns: _Writer, self