code: def gather_keys(conf: Configs):
      """
      ## Gather $\big(f(c_i), w_i\big)$ and save them in numpy arrays
  
      *Note that these numpy arrays will take up a lot of space (even few hundred gigabytes)
      depending on the size of your dataset*.
      """
  
      # 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
      # Numpy array for $f(c_i)$
      keys_store = np.memmap(
          str(lab.get_data_path() / "keys.npy"),
          dtype=np.float32,
          mode="w+",
          shape=(n_keys, d_model),
      )
      # Numpy array for $w_i$
      vals_store = np.memmap(
          str(lab.get_data_path() / "vals.npy"),
          dtype=np.int,
          mode="w+",
          shape=(n_keys, 1),
      )
  
      # Number of keys $f(c_i)$ collected
      added = 0
      with torch.no_grad():
          # Loop through data
          for i, batch in monit.enum(
              "Collect data", data_loader, is_children_silent=True
          ):
              # $w_i$ the target labels
              vals = batch[1].view(-1, 1)
              # Input data moved to the device of the model
              data = batch[0].to(conf.device)
              # Run the model
              _ = conf.model(data)
              # Get $f(c_i)$
              keys = conf.model.ff_input.view(-1, d_model)
              # Save keys, $f(c_i)$ in the memory mapped numpy array
              keys_store[added : added + keys.shape[0]] = keys.cpu()
              # Save values, $w_i$ in the memory mapped numpy array
              vals_store[added : added + keys.shape[0]] = vals
              # Increment the number of collected keys
              added += keys.shape[0]
  
file_path: annotated_deep_learning_paper_implementations/labml_nn/transformers/knn/build_index.py
imports: from typing import Optional
  import faiss
  import numpy
  import torch
  from labml import experiment, monit, lab
  from labml.utils.pytorch import get_modules
  from labml_nn.transformers.knn.train_model import Configs
label_desc: Use image node_id 2 for calling a global function with example usage: gather_keys(conf) without return types