code: def main():
      parser = ArgumentParser()
      parser.add_argument(
          "--lm_model_file",
          type=str,
          required=True,
          help="path to LM model .nemo file or the name of a HuggingFace pretrained models like 'transfo-xl-wt103' or 'gpt2'",
      )
      parser.add_argument(
          "--beams_file",
          type=str,
          required=True,
          help="path to beams .tsv file",
      )
      parser.add_argument(
          "--eval_manifest",
          type=str,
          required=True,
          help="path to the evaluation `.json` manifest file",
      )
      parser.add_argument(
          "--beam_size",
          type=int,
          required=True,
          help="number of beams per candidate",
      )
      parser.add_argument(
          "--batch_size",
          type=int,
          default=256,
          help="inference batch size",
      )
      parser.add_argument(
          "--alpha",
          type=float,
          default=None,
          help="parameter alpha of the fusion",
      )
      parser.add_argument(
          "--beta",
          type=float,
          default=None,
          help="parameter beta of the fusion",
      )
      parser.add_argument(
          "--max_seq_length",
          default=512,
          help="Maximum sequence length (in tokens) for the input",
      )
      parser.add_argument(
          "--scores_output_file",
          default=None,
          type=str,
          help="The optional path to store the rescored beams",
      )
      parser.add_argument(
          "--device",
          default="cuda",
          type=str,
          help="The device to load the model onto to calculate the scores",
      )
      parser.add_argument(
          "--use_amp",
          action="store_true",
          help="Whether to use AMP if available to calculate the scores",
      )
      args = parser.parse_args()
  
      device = args.device
      if device.startswith("cuda") and not torch.cuda.is_available():
          logging.info(f"cuda is not available! switched to cpu.")
          device = "cpu"
  
      if args.lm_model_file.endswith(".nemo"):
          nemo_model = True
          logging.info("Attempting to initialize from .nemo file...")
          model = TransformerLMModel.restore_from(
              restore_path=args.lm_model_file,
              map_location=torch.device(device),
          ).eval()
          model_tokenizer = model.tokenizer
      else:
          nemo_model = False
          logging.info(
              "Attempting to initialize from a pretrained model from HuggingFace..."
          )
          model = (
              AutoModelForCausalLM.from_pretrained(
                  pretrained_model_name_or_path=args.lm_model_file,
                  is_decoder=True,
              )
              .to(device)
              .eval()
          )
          model_tokenizer = get_tokenizer(
              tokenizer_name=args.lm_model_file
          )
  
      max_seq_length = args.max_seq_length
      dataset = BeamScoresDataset(
          args.beams_file,
          model_tokenizer,
          args.eval_manifest,
          args.beam_size,
          max_seq_length,
      )
      data_loader = torch.utils.data.DataLoader(
          dataset=dataset, batch_size=args.batch_size
      )
  
      if args.use_amp:
          if (
              torch.cuda.is_available()
              and hasattr(torch.cuda, "amp")
              and hasattr(torch.cuda.amp, "autocast")
          ):
              logging.info("AMP is enabled!\n")
              autocast = torch.cuda.amp.autocast
      else:
  
          @contextlib.contextmanager
          def autocast():
              yield
  
      if "attention_mask" in inspect.getfullargspec(model.forward).args:
          support_att_mask = True
      else:
          support_att_mask = False
      logging.info(f"Rescoring with beam_size: {args.beam_size}")
      logging.info("Calculating the scores...")
      with autocast():
          with torch.no_grad():
              am_scores, lm_scores, dists, ref_lens, lens_in_chars = (
                  [],
                  [],
                  [],
                  [],
                  [],
              )
              for batch in tqdm.tqdm(data_loader):
                  (
                      input_ids,
                      input_mask,
                      acoustic_score,
                      dist,
                      ref_len,
                      len_in_chars,
                      idx,
                  ) = batch
  
                  max_len_in_batch = (
                      input_mask.sum(dim=0).argmin().item()
                  )
                  input_ids, input_mask = (
                      input_ids[:, :max_len_in_batch],
                      input_mask[:, :max_len_in_batch],
                  )
                  if torch.cuda.is_available():
                      input_ids, input_mask = input_ids.to(
                          device
                      ), input_mask.to(device)
                      dist, acoustic_score, len_in_chars = (
                          dist.to(device),
                          acoustic_score.to(device),
                          len_in_chars.to(device),
                      )
                  # some models like Transformer-XL don't need attention_mask as input
                  if support_att_mask:
                      log_probs = model(
                          input_ids=input_ids, attention_mask=input_mask
                      )
                  else:
                      log_probs = model(input_ids=input_ids)
  
                  if not nemo_model:
                      log_probs = torch.nn.functional.log_softmax(
                          log_probs.logits, dim=-1
                      )
  
                  target_log_probs = (
                      log_probs[:, :-1]
                      .gather(2, input_ids[:, 1:].unsqueeze(2))
                      .squeeze(2)
                  )
                  neural_lm_score = torch.sum(
                      target_log_probs * input_mask[:, 1:], dim=-1
                  )
  
                  am_scores.append(acoustic_score)
                  lm_scores.append(neural_lm_score)
                  dists.append(dist)
                  ref_lens.append(ref_len)
                  lens_in_chars.append(len_in_chars)
  
      am_scores = torch.cat(am_scores).view(-1, args.beam_size)
      lm_scores = torch.cat(lm_scores).view(-1, args.beam_size)
      dists = torch.cat(dists).view(-1, args.beam_size)
      ref_lens = torch.cat(ref_lens).view(-1, args.beam_size)
      lens_in_chars = (
          torch.cat(lens_in_chars)
          .view(-1, args.beam_size)
          .to(am_scores.dtype)
      )
  
      total_len = ref_lens[:, 0].sum()
      model_wer = dists[:, 0].sum() / total_len
      ideal_wer = dists.min(dim=1)[0].sum() / total_len
  
      if args.alpha is None:
          logging.info("Linear search for alpha...")
          coef1, _ = linear_search_wer(
              dists=dists,
              scores1=am_scores,
              scores2=lm_scores,
              total_len=total_len,
              param_name="alpha",
          )
          coef1 = np.round(coef1, 3)
          logging.info(f"alpha={coef1} achieved the best WER.")
          logging.info(
              f"------------------------------------------------"
          )
      else:
          coef1 = args.alpha
  
      scores = am_scores + coef1 * lm_scores
  
      if args.beta is None:
          logging.info("Linear search for beta...")
          coef2, _ = linear_search_wer(
              dists=dists,
              scores1=scores,
              scores2=lens_in_chars,
              total_len=total_len,
              param_name="beta",
          )
          coef2 = np.round(coef2, 3)
          logging.info(f"beta={coef2} achieved the best WER.")
          logging.info(
              f"------------------------------------------------"
          )
      else:
          coef2 = args.beta
  
      new_scores = am_scores + coef1 * lm_scores + coef2 * lens_in_chars
      rescored_wer = compute_wer(dists, new_scores, total_len)
  
      logging.info(
          f"Input beams WER: {np.round(model_wer.item() * 100, 2)}%"
      )
      logging.info(f"------------------------------------------------")
      logging.info(
          f"  +LM rescoring WER: {np.round(rescored_wer * 100, 2)}%"
      )
      logging.info(f"  with alpha={coef1}, beta={coef2}")
      logging.info(f"------------------------------------------------")
      logging.info(
          f"Oracle WER: {np.round(ideal_wer.item() * 100, 2)}%"
      )
      logging.info(f"------------------------------------------------")
  
      new_scores_flatten = new_scores.flatten()
      if args.scores_output_file is not None:
          logging.info(
              f"Saving the candidates with their new scores at `{args.scores_output_file}`..."
          )
          with open(
              args.scores_output_file, "w", encoding="utf-8"
          ) as fout:
              for sample_id in range(len(dataset)):
                  fout.write(
                      f"{dataset.data[0][sample_id]}\t{new_scores_flatten[sample_id]}\n"
                  )
  
imports: import contextlib
  import inspect
  import json
  from argparse import ArgumentParser
  import editdistance
  import matplotlib.pyplot
  import numpy
  import pandas
  import torch
  import tqdm
  from transformers import AutoModelForCausalLM
  from nemo.collections.nlp.models.language_modeling import TransformerLMModel
  from nemo.collections.nlp.modules.common.tokenizer_utils import get_tokenizer
  from nemo.utils import logging
label_desc: Use image node_id 3 for calling a global function with example usage: main() without return types