code: def main():
      parser = ArgumentParser()
      parser.add_argument(
          "--reverse_model",
          type=str,
          help="Path to .nemo model file(s). If ensembling, provide comma separated paths to multiple models.",
      )
      parser.add_argument(
          "--language_model",
          type=str,
          help="Optional path to an LM model that has the same tokenizer as NMT models.",
      )
      parser.add_argument(
          "--forward_model_coef",
          type=float,
          default=1.0,
          help="Weight assigned to the forward NMT model for re-ranking.",
      )
      parser.add_argument(
          "--reverse_model_coef",
          type=float,
          default=0.7,
          help="Weight assigned to the reverse NMT model for re-ranking.",
      )
      parser.add_argument(
          "--target_lm_coef",
          type=float,
          default=0.07,
          help="Weight assigned to the target LM model for re-ranking.",
      )
      parser.add_argument(
          "--srctext",
          type=str,
          default=None,
          help="Path to a TSV file containing forward model scores of the format source \t beam_candidate_i \t forward_score",
      )
      parser.add_argument(
          "--cached_score_file",
          type=str,
          default=None,
          help="Path to a TSV file containing cached scores for each beam candidate. Format source \t target \t forward_score \t reverse_score \t lm_score \t src_len \t tgt_len",
      )
      parser.add_argument(
          "--tgtout",
          type=str,
          required=True,
          help="Path to the file where re-ranked translations are to be written.",
      )
      parser.add_argument(
          "--beam_size",
          type=int,
          default=4,
          help="Beam size with which forward model translations were generated. IMPORTANT: mismatch can lead to wrong results and an incorrect number of generated translations.",
      )
      parser.add_argument(
          "--target_lang",
          type=str,
          default=None,
          help="Target language identifier ex: en,de,fr,es etc.",
      )
      parser.add_argument(
          "--source_lang",
          type=str,
          default=None,
          help="Source language identifier ex: en,de,fr,es etc.",
      )
      parser.add_argument(
          "--write_scores",
          action="store_true",
          help="Whether to write forward, reverse and lm scores to a file.",
      )
      parser.add_argument(
          "--length_normalize_scores",
          action="store_true",
          help="If true, it will divide forward, reverse and lm scores by the corresponding sequence length.",
      )
      parser.add_argument(
          "--len_pen",
          type=float,
          default=None,
          help="Apply a length penalty based on target lengths to the final NCR score.",
      )
  
      args = parser.parse_args()
      torch.set_grad_enabled(False)
  
      if args.cached_score_file is None:
          reverse_models = []
          for model_path in args.reverse_model.split(","):
              if not model_path.endswith(".nemo"):
                  raise NotImplementedError(
                      f"Only support .nemo files, but got: {model_path}"
                  )
              model = nemo_nlp.models.machine_translation.MTEncDecModel.restore_from(
                  restore_path=model_path
              ).eval()
              model.eval_loss_fn.reduction = "none"
              reverse_models.append(model)
  
          lm_model = nemo_nlp.models.language_modeling.TransformerLMModel.restore_from(
              restore_path=args.language_model
          ).eval()
  
      if (
          args.srctext is not None
          and args.cached_score_file is not None
      ):
          raise ValueError(
              "Only one of --srctext or --cached_score_file must be provided."
          )
  
      if args.srctext is None and args.cached_score_file is None:
          raise ValueError(
              "Neither --srctext nor --cached_score_file were provided."
          )
  
      if args.srctext is not None:
          logging.info(f"Re-ranking: {args.srctext}")
      else:
          logging.info(
              f"Re-ranking from cached score file only: {args.cached_score_file}"
          )
  
      if args.cached_score_file is None:
          if torch.cuda.is_available():
              reverse_models = [
                  model.cuda() for model in reverse_models
              ]
              lm_model = lm_model.cuda()
  
      src_text = []
      tgt_text = []
      all_reverse_scores = []
      all_lm_scores = []
      all_forward_scores = []
      all_src_lens = []
      all_tgt_lens = []
  
      # Chceck args if re-ranking from cached score file.
      if args.cached_score_file is not None:
          if args.write_scores:
              raise ValueError(
                  "--write_scores cannot be provided with a cached score file."
              )
          if args.reverse_model is not None:
              raise ValueError(
                  "--reverse_model cannot be provided with a cached score file since it assumes reverse scores already present in the cached file."
              )
          if args.language_model is not None:
              raise ValueError(
                  "--language_model cannot be provided with a cached score file since it assumes language model scores already present in the cached file."
              )
  
      if args.srctext is not None:
          # Compute reverse scores and LM scores from the provided models since cached scores file is not provided.
          with open(args.srctext, "r") as src_f:
              count = 0
              for line in src_f:
                  src_text.append(line.strip().split("\t"))
                  if len(src_text) == args.beam_size:
                      # Source and target sequences are flipped for the reverse direction model.
                      src_texts = [item[1] for item in src_text]
                      tgt_texts = [item[0] for item in src_text]
                      src, src_mask = reverse_models[
                          0
                      ].prepare_inference_batch(src_texts)
                      tgt, tgt_mask = reverse_models[
                          0
                      ].prepare_inference_batch(tgt_texts, target=True)
                      src_lens = src_mask.sum(1).data.cpu().tolist()
                      tgt_lens = tgt_mask.sum(1).data.cpu().tolist()
                      forward_scores = [
                          float(item[2]) for item in src_text
                      ]
  
                      # Ensemble of reverse model scores.
                      nmt_lls = []
                      for model in reverse_models:
                          nmt_log_probs = model(
                              src,
                              src_mask,
                              tgt[:, :-1],
                              tgt_mask[:, :-1],
                          )
                          nmt_nll = model.eval_loss_fn(
                              log_probs=nmt_log_probs, labels=tgt[:, 1:]
                          )
                          nmt_ll = (
                              nmt_nll.view(
                                  nmt_log_probs.size(0),
                                  nmt_log_probs.size(1),
                              ).sum(1)
                              * -1.0
                          )
                          nmt_ll = nmt_ll.data.cpu().numpy().tolist()
                          nmt_lls.append(nmt_ll)
                      reverse_scores = np.stack(nmt_lls).mean(0)
  
                      # LM scores.
                      if lm_model is not None:
                          # Compute LM score for the src of the reverse model.
                          lm_log_probs = lm_model(
                              src[:, :-1], src_mask[:, :-1]
                          )
                          lm_nll = model.eval_loss_fn(
                              log_probs=lm_log_probs, labels=src[:, 1:]
                          )
                          lm_ll = (
                              lm_nll.view(
                                  lm_log_probs.size(0),
                                  lm_log_probs.size(1),
                              ).sum(1)
                              * -1.0
                          )
                          lm_ll = lm_ll.data.cpu().numpy().tolist()
                      else:
                          lm_ll = None
                      lm_scores = lm_ll
  
                      all_reverse_scores.extend(reverse_scores)
                      all_lm_scores.extend(lm_scores)
                      all_forward_scores.extend(forward_scores)
  
                      # Swapping source and target here back again since this is what gets written to the file.
                      all_src_lens.extend(tgt_lens)
                      all_tgt_lens.extend(src_lens)
  
                      fused_scores = score_fusion(
                          args,
                          forward_scores,
                          reverse_scores,
                          lm_scores,
                          src_lens,
                          tgt_lens,
                      )
                      tgt_text.append(
                          src_texts[np.argmax(fused_scores)]
                      )
                      src_text = []
                      count += 1
                      print(f"Reranked {count} sentences")
  
      else:
          # Use reverse and LM scores from the cached scores file to re-rank.
          with open(args.cached_score_file, "r") as src_f:
              count = 0
              for line in src_f:
                  src_text.append(line.strip().split("\t"))
                  if len(src_text) == args.beam_size:
                      if not all([len(item) == 7 for item in src_text]):
                          raise IndexError(
                              "All lines did not contain exactly 5 fields. Format - src_txt \t tgt_text \t forward_score \t reverse_score \t lm_score \t src_len \t tgt_len"
                          )
                      src_texts = [item[0] for item in src_text]
                      tgt_texts = [item[1] for item in src_text]
                      forward_scores = [
                          float(item[2]) for item in src_text
                      ]
                      reverse_scores = [
                          float(item[3]) for item in src_text
                      ]
                      lm_scores = [float(item[4]) for item in src_text]
                      src_lens = [float(item[5]) for item in src_text]
                      tgt_lens = [float(item[6]) for item in src_text]
  
                      fused_scores = score_fusion(
                          args,
                          forward_scores,
                          reverse_scores,
                          lm_scores,
                          src_lens,
                          tgt_lens,
                      )
                      tgt_text.append(
                          tgt_texts[np.argmax(fused_scores)]
                      )
                      src_text = []
                      count += 1
                      print(f"Reranked {count} sentences")
  
      with open(args.tgtout, "w") as tgt_f:
          for line in tgt_text:
              tgt_f.write(line + "\n")
  
      # Write scores file
      if args.write_scores:
          with open(args.tgtout + ".scores", "w") as tgt_f, open(
              args.srctext, "r"
          ) as src_f:
              src_lines = []
              for line in src_f:
                  src_lines.append(line.strip().split("\t"))
              if not (
                  len(all_reverse_scores)
                  == len(all_lm_scores)
                  == len(all_forward_scores)
                  == len(src_lines)
              ):
                  raise ValueError(
                      f"Length of scores files do not match. {len(all_reverse_scores)} != {len(all_lm_scores)} != {len(all_forward_scores)} != {len(src_lines)}. This is most likely because --beam_size is set incorrectly. This needs to be set to the same value that was used to generate translations."
                  )
              for f, r, lm, sl, tl, src in zip(
                  all_forward_scores,
                  all_reverse_scores,
                  all_lm_scores,
                  all_src_lens,
                  all_tgt_lens,
                  src_lines,
              ):
                  tgt_f.write(
                      src[0]
                      + "\t"
                      + src[1]
                      + "\t"
                      + str(f)
                      + "\t"
                      + str(r)
                      + "\t"
                      + str(lm)
                      + "\t"
                      + str(sl)
                      + "\t"
                      + str(tl)
                      + "\n"
                  )
  
imports: from argparse import ArgumentParser
  import numpy
  import torch
  import nemo.collections.nlp
  from nemo.utils import logging
label_desc: Use image node_id 2 for calling a global function with example usage: main() without return types