code: def compute_top_k(
      task,
      models,
      dataset,
      k,
      use_cuda,
      max_tokens=None,
      max_sentences=None,
      progress_bar_args=None,
  ):
      """
      This function runs forward computation on an ensemble of trained models
      using binarized parallel training data and returns the top-k probabilities
      and their corresponding token indices for each output step.
  
      Returns: (top_k_scores, top_k_indices)
          Each a NumPy array of size (total_target_tokens, k)
      """
      top_k_scores_list = [None for _ in range(len(dataset))]
      top_k_indices_list = [None for _ in range(len(dataset))]
  
      itr = task.get_batch_iterator(
          dataset=dataset,
          max_tokens=max_tokens,
          max_sentences=max_sentences,
      ).next_epoch_itr(shuffle=False)
      if progress_bar_args is not None:
          itr = progress_bar.build_progress_bar(
              args=progress_bar_args,
              iterator=itr,
              prefix=f"top-k probs eval",
              no_progress_bar="simple",
          )
  
      for sample in itr:
          sentence_ids = sample["id"]
          target_lengths = (
              (
                  sample["net_input"]["prev_output_tokens"]
                  != dataset.tgt_dict.pad()
              )
              .sum(axis=1)
              .numpy()
          )
          if use_cuda:
              sample = utils.move_to_cuda(sample)
          avg_probs = None
          for model in models:
              with torch.no_grad():
                  net_output = model(**sample["net_input"])
                  probs = model.get_normalized_probs(
                      net_output, log_probs=False
                  )
              if avg_probs is None:
                  avg_probs = probs
              else:
                  avg_probs.add_(probs)
          avg_probs.div_(len(models))
  
          top_k_avg_probs, indices = torch.topk(avg_probs, k=k)
  
          top_k_probs_normalized = F.normalize(
              top_k_avg_probs, p=1, dim=2
          ).cpu()
          indices = indices.cpu()
  
          for i, sentence_id in enumerate(sentence_ids):
              length = target_lengths[i]
              top_k_scores_list[sentence_id] = top_k_probs_normalized[
                  i
              ][:length].numpy()
              top_k_indices_list[sentence_id] = indices[i][
                  :length
              ].numpy()
  
      assert all(
          top_k_scores is not None for top_k_scores in top_k_scores_list
      ), "scores not calculated for all examples!"
      assert all(
          top_k_indices is not None
          for top_k_indices in top_k_indices_list
      ), "indices not calculated for all examples!"
  
      top_k_scores = np.concatenate(top_k_scores_list, axis=0)
      top_k_indices = np.concatenate(top_k_indices_list, axis=0)
  
      return top_k_scores, top_k_indices
  
file_path: translate/pytorch_translate/research/knowledge_distillation/collect_top_k_probs.py
imports: import numpy
  import torch
  import torch.nn.functional
  from fairseq import options, progress_bar, utils
  from pytorch_translate import hybrid_transformer_rnn, options, rnn, transformer, utils
  from pytorch_translate.constants import CHECKPOINT_PATHS_DELIMITER
  from pytorch_translate.tasks.pytorch_translate_multi_task import PyTorchTranslateMultiTask
label_desc: Use image node_id 1 for calling a global function with example usage: compute_top_k(task, models, dataset, k, use_cuda, max_tokens, max_sentences, progress_bar_args) and returns: top_k_scores,  top_k_indices