code: def save_top_k(args):
      """
      This function runs forward computation on an ensemble of trained models
      using binarized parallel training data and saves the top-k probabilities
      and their corresponding token indices for each output step.
  
      Note that the Python binary accepts all generation params, but ignores
      inapplicable ones (such as those related to output length). --max-tokens
      is of particular importance to prevent memory errors.
      """
      pytorch_translate_options.print_args(args)
      use_cuda = torch.cuda.is_available() and not getattr(
          args, "cpu", False
      )
  
      (
          models,
          model_args,
          task,
      ) = pytorch_translate_utils.load_diverse_ensemble_for_inference(
          args.path.split(CHECKPOINT_PATHS_DELIMITER)
      )
      for model in models:
          model.eval()
          if use_cuda:
              model.cuda()
  
      append_eos_to_source = model_args[0].append_eos_to_source
      reverse_source = model_args[0].reverse_source
      assert all(
          a.append_eos_to_source == append_eos_to_source
          and a.reverse_source == reverse_source
          for a in model_args
      )
      assert (
          args.source_binary_file != ""
          and args.target_binary_file != ""
      ), "collect_top_k_probs requires binarized data."
      task.load_dataset(
          args.gen_subset,
          args.source_binary_file,
          args.target_binary_file,
      )
  
      assert (
          args.top_k_probs_binary_file != ""
      ), "must specify output file (--top-k-probs-binary-file)!"
      output_path = args.top_k_probs_binary_file
  
      dataset = task.dataset(args.gen_subset)
  
      top_k_scores, top_k_indices = compute_top_k(
          task=task,
          models=models,
          dataset=dataset,
          k=args.k_probs_to_collect,
          use_cuda=use_cuda,
          max_tokens=args.teacher_max_tokens,
          max_sentences=args.max_sentences,
          progress_bar_args=args,
      )
  
      np.savez(
          output_path,
          top_k_scores=top_k_scores,
          top_k_indices=top_k_indices,
      )
      print(
          f"Saved top {top_k_scores.shape[1]} probs for a total of "
          f"{top_k_scores.shape[0]} tokens to file {output_path}"
      )
  
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 2 for calling a global function with example usage: save_top_k(args) without return types