code: def build_index(
      conf: Configs,
      n_centeroids: int = 2048,
      code_size: int = 64,
      n_probe: int = 8,
      n_train: int = 200_000,
  ):
      """
      ## Build FAISS index
  
      [Getting started](https://github.com/facebookresearch/faiss/wiki/Getting-started),
      [faster search](https://github.com/facebookresearch/faiss/wiki/Faster-search),
      and [lower memory footprint](https://github.com/facebookresearch/faiss/wiki/Lower-memory-footprint)
      tutorials on FAISS will help you learn more about FAISS usage.
      """
      # 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
  
      # Build an index with Verenoi cell based faster search with compression that
      # doesn't store full vectors.
      quantizer = faiss.IndexFlatL2(d_model)
      index = faiss.IndexIVFPQ(
          quantizer, d_model, n_centeroids, code_size, 8
      )
      index.nprobe = n_probe
  
      # Load the memory mapped numpy array of keys
      keys_store = np.memmap(
          str(lab.get_data_path() / "keys.npy"),
          dtype=np.float32,
          mode="r",
          shape=(n_keys, d_model),
      )
  
      # Pick a random sample of keys to train the index with
      random_sample = np.random.choice(
          np.arange(n_keys), size=[min(n_train, n_keys)], replace=False
      )
  
      with monit.section("Train index"):
          # Train the index to store the keys
          index.train(keys_store[random_sample])
  
      # Add keys to the index; $\big(f(c_i), i\big)$
      for s in monit.iterate("Index", range(0, n_keys, 1024)):
          e = min(s + 1024, n_keys)
          # $f(c_i)$
          keys = keys_store[s:e]
          # $i$
          idx = np.arange(s, e)
          # Add to index
          index.add_with_ids(keys, idx)
  
      with monit.section("Save"):
          # Save the index
          faiss.write_index(
              index, str(lab.get_data_path() / "faiss.index")
          )
  
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 3 for calling a global function with example usage: build_index(conf, n_centeroids, code_size, n_probe, n_train) without return types