code: def load_index(conf: Configs, n_probe: int = 8):
      """
      ## Load the index
      """
      # Dimensions of $f(c_i)$
      d_model = conf.transformer.d_model
      # Training data loader
      data_loader = conf.trainer.data_loader
      # Number of contexts; i.e. number of tokens in the training data minus one.
      # $\big(f(c_i), w_i\big)$ for $i \in [2, T]$
      n_keys = data_loader.data.shape[0] * data_loader.data.shape[1] - 1
  
      # Load FAISS index
      with monit.section("Load index"):
          index = faiss.read_index(
              str(lab.get_data_path() / "faiss.index")
          )
      # Set number of cells to probe
      index.nprobe = n_probe
  
      # Load memory mapped numpy arrays
      keys_store = np.memmap(
          str(lab.get_data_path() / "keys.npy"),
          dtype=np.float32,
          mode="r",
          shape=(n_keys, d_model),
      )
      vals_store = np.memmap(
          str(lab.get_data_path() / "vals.npy"),
          dtype=np.int,
          mode="r",
          shape=(n_keys, 1),
      )
  
      return index, keys_store, vals_store
  
file_path: annotated_deep_learning_paper_implementations/labml_nn/transformers/knn/eval_knn.py
imports: from typing import Optional, List
  import faiss
  import numpy
  import torch
  from labml import monit, lab
  from labml.logger import inspect
  from labml_nn.transformers.knn.train_model import Configs
label_desc: Use image node_id 3 for calling a global function with example usage: load_index(conf, n_probe) and returns: index,  keys_store,  vals_store