code: def validation_loss(
      knn_weights: List[float],
      last_n: Optional[int],
      conf: Configs,
      index: faiss.IndexFlatL2,
      keys_store: np.ndarray,
      vals_store: np.ndarray,
  ):
      """
      ## Calculate validation loss
  
      We calculate the validation loss of the combined on $k$-NN prediction and transformer prediction.
      The weight given to the $k$-NN model is given by `knn_weight`.
      It's a list of weights and we calculate the validation loss for each.
      """
  
      # List of losses for each `knn_weights`
      losses = [[] for _ in knn_weights]
      # Number of samples in each batch
      n_samples = []
      with torch.no_grad():
          # Iterate through validation data
          for i, batch in monit.enum(
              "Validation",
              conf.validator.data_loader,
              is_children_silent=True,
          ):
              # Get data and target labels
              data, target = batch[0].to(conf.device), batch[1].to(
                  conf.device
              )
              # Run the model and get predictions $p(w_t, c_t)$
              res = conf.model(data)
              # Get $k$-NN predictions
              res_knn = knn(
                  conf.model.ff_input.cpu(),
                  index,
                  keys_store,
                  vals_store,
                  conf.n_tokens,
              )
              res_knn = res_knn.to(conf.device)
  
              # This is to calculate only the loss for `last_n` tokens.
              # This is important because the first predictions (along the sequence)
              # of transformer model has very few past tokens to look at.
              if last_n:
                  res = res[-last_n:]
                  res_knn = res_knn[-last_n:]
                  target = target[-last_n:]
  
              # Number of samples
              n_s = res.shape[0] * data.shape[1]
              n_samples.append(n_s)
  
              # Calculate scores for each of `knn_weights`.
              for i, c in enumerate(knn_weights):
                  # Calculate the loss
                  loss = conf.loss_func(
                      res_knn * c + (1 - c) * res, target
                  )
                  losses[i].append(loss * n_s)
  
      return losses, n_samples
  
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 2 for calling a global function with example usage: validation_loss(knn_weights, last_n, conf, index, keys_store, vals_store) and returns: losses,  n_samples