class_docstr: https://arxiv.org/abs/2306.02317
  BERT-based model for Spellchecking ASR Customization.
  It takes as input ASR hypothesis and candidate customization entries.
  It labels the hypothesis with correct entry index or 0.
  Example input:   [CLS] a s t r o n o m e r s _ d i d i e _ s o m o n _ a n d _ t r i s t i a n _ g l l o [SEP] d i d i e r _ s a u m o n [SEP] a s t r o n o m i e [SEP] t r i s t a n _ g u i l l o t [SEP] ...
  Input segments:      0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0     1 1 1 1 1 1 1 1 1 1 1 1 1 1     2 2 2 2 2 2 2 2 2 2 2     3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3     4      
  Example output:      0 0 0 0 0 0 0 0 0 0 0 0 0 1 1 1 1 1 1 1 1 1 1 1 0 0 0 0 0 3 3 3 3 3 3 3 3 3 3 3 3 3 0     ...
class_docstr_tok: ["https",":","\/\/arxiv.org\/abs\/2306.02317","BERT-based","model","for","Spellchecking","ASR","Customization",".","It","takes","as","input","ASR","hypothesis","and","candidate","customization","entries",".","It","labels","the","hypothesis","with","correct","entry","index","or","0",".","Example","input",":","[","CLS","]","a","s","t","r","o","n","o","m","e","r","s","_","d","i","d","i","e","_","s","o","m","o","n","_","a","n","d","_","t","r","i","s","t","i","a","n","_","g","l","l","o","[","SEP","]","d","i","d","i","e","r","_","s","a","u","m","o","n","[","SEP","]","a","s","t","r","o","n","o","m","i","e","[","SEP","]","t","r","i","s","t","a","n","_","g","u","i","l","l","o","t","[","SEP","]","...","Input","segments",":","0","0","0","0","0","0","0","0","0","0","0","0","0","0","0","0","0","0","0","0","0","0","0","0","0","0","0","0","0","0","0","0","0","0","0","0","0","0","0","0","0","0","1","1","1","1","1","1","1","1","1","1","1","1","1","1","2","2","2","2","2","2","2","2","2","2","2","3","3","3","3","3","3","3","3","3","3","3","3","3","3","3","3","4","Example","output",":","0","0","0","0","0","0","0","0","0","0","0","0","0","1","1","1","1","1","1","1","1","1","1","1","0","0","0","0","0","3","3","3","3","3","3","3","3","3","3","3","3","3","0","..."]
docstr: Main function for Inference
  
  Args:
      dataloader_cfg: config for dataloader
      input_name: Input file with tab-separated text records. Each record consists of 2 items:
          - ASR hypothesis
          - candidate phrases separated by semicolon
      output_name: Output file with tab-separated text records. Each record consists of 2 items:
          - ASR hypothesis
          - candidate phrases separated by semicolon
          - list of possible replacements with probabilities (start, pos, candidate_id, prob), separated by semicolon
          - list of labels, predicted for each letter (for debug purposes)
  
  Returns: None
docstr_tok: ["Main","function","for","Inference","Args",":","dataloader_cfg",":","config","for","dataloader","input_name",":","Input","file","with","tab-separated","text","records",".","Each","record","consists","of","2","items",":","-","ASR","hypothesis","-","candidate","phrases","separated","by","semicolon","output_name",":","Output","file","with","tab-separated","text","records",".","Each","record","consists","of","2","items",":","-","ASR","hypothesis","-","candidate","phrases","separated","by","semicolon","-","list","of","possible","replacements","with","probabilities","(","start",",","pos",",","candidate_id",",","prob",")",",","separated","by","semicolon","-","list","of","labels",",","predicted","for","each","letter","(","for","debug","purposes",")","Returns",":","None"]
code: def infer(
      self,
      dataloader_cfg: DictConfig,
      input_name: str,
      output_name: str,
  ) -> None:
      """Main function for Inference
  
      Args:
          dataloader_cfg: config for dataloader
          input_name: Input file with tab-separated text records. Each record consists of 2 items:
              - ASR hypothesis
              - candidate phrases separated by semicolon
          output_name: Output file with tab-separated text records. Each record consists of 2 items:
              - ASR hypothesis
              - candidate phrases separated by semicolon
              - list of possible replacements with probabilities (start, pos, candidate_id, prob), separated by semicolon
              - list of labels, predicted for each letter (for debug purposes)
  
      Returns: None
      """
      mode = self.training
      device = "cuda" if torch.cuda.is_available() else "cpu"
  
      try:
          # Switch model to evaluation mode
          self.eval()
          self.to(device)
          logging_level = logging.get_verbosity()
          logging.set_verbosity(logging.WARNING)
          infer_datalayer = self._setup_infer_dataloader(
              dataloader_cfg, input_name
          )
  
          all_tag_preds = (
              []
          )  # list(size=number of sentences) of lists(size=number of letters) of tag predictions (best candidate_id for each letter)
          all_possible_replacements = (
              []
          )  # list(size=number of sentences) of lists(size=number of potential replacements) of tuples(start, pos, candidate_id, prob)
          for batch in iter(infer_datalayer):
              (
                  input_ids,
                  input_mask,
                  segment_ids,
                  input_ids_for_subwords,
                  input_mask_for_subwords,
                  segment_ids_for_subwords,
                  character_pos_to_subword_pos,
                  fragment_indices,
              ) = batch
  
              # tag_logits.shape = [batch_size, char_seq_len, num_labels]; num_labels=11: ids from 0 to 10; .dtype=float32
              tag_logits = self.forward(
                  input_ids=input_ids.to(self.device),
                  input_mask=input_mask.to(self.device),
                  segment_ids=segment_ids.to(self.device),
                  input_ids_for_subwords=input_ids_for_subwords.to(
                      self.device
                  ),
                  input_mask_for_subwords=input_mask_for_subwords.to(
                      self.device
                  ),
                  segment_ids_for_subwords=segment_ids_for_subwords.to(
                      self.device
                  ),
                  character_pos_to_subword_pos=character_pos_to_subword_pos.to(
                      self.device
                  ),
              )
  
              # fragment_indices.shape=[batsh_size, num_fragments, 3], where last dimension is [start, end, label], where label is candidate id from 1 to 10
              # Next we want to convert predictions for separate letters to probabilities for each whole fragment from fragment_indices.
              # To achieve this we first sum the letter logits in each fragment and divide by its length.
              # (We use .cumsum and then difference between end and start to get sum per fragment).
              # Then we convert logits to probs with softmax and for each fragment extract only the prob for given label.
              # Finally we get a list of tuples (start, end, label, prob)
              indices_len = fragment_indices.shape[1]
              # this padding adds a row of zeros (size=num_labels) as first element of sequence in second dimension. This is needed for cumsum operations.
              padded_logits = torch.nn.functional.pad(
                  tag_logits, pad=(0, 0, 1, 0)
              )
              (
                  batch_size,
                  seq_len,
                  num_labels,
              ) = (
                  padded_logits.shape
              )  # seq_len is +1 compared to that of tag_logits, because of padding
              # cumsum.shape=[batch_size, seq_len, num_labels]
              cumsum = padded_logits.cumsum(dim=1)
              # the size -1 is inferred from other dimensions. We get rid of batch dimension.
              cumsum_view = cumsum.view(-1, num_labels)
              word_index = (
                  torch.ones(
                      (batch_size, indices_len), dtype=torch.long
                  )
                  * torch.arange(batch_size).reshape((-1, 1))
                  * seq_len
              ).view(-1)
              lower_index = (fragment_indices[..., 0]).view(
                  -1
              ) + word_index
              higher_index = (fragment_indices[..., 1]).view(
                  -1
              ) + word_index
              d_index = (
                  (higher_index - lower_index)
                  .reshape((-1, 1))
                  .to(self.device)
              )  # word lengths
              dlog = (
                  cumsum_view[higher_index, :]
                  - cumsum_view[lower_index, :]
              )  # sum of logits
              # word_logits.shape=[batch_size, indices_len, num_labels]
              word_logits = (dlog / d_index.float()).view(
                  batch_size, indices_len, num_labels
              )
              # convert logits to probs, same shape
              word_probs = torch.nn.functional.softmax(
                  word_logits, dim=-1
              ).to(self.device)
              # candidate_index.shape=[batch_size, indices_len]
              candidate_index = fragment_indices[:, :, 2].to(
                  self.device
              )
              # candidate_probs.shape=[batch_size, indices_len]
              candidate_probs = torch.take_along_dim(
                  word_probs, candidate_index.unsqueeze(2), dim=-1
              ).squeeze(2)
              for i in range(batch_size):
                  possible_replacements = []
                  for j in range(indices_len):
                      start, end, candidate_id = (
                          int(fragment_indices[i][j][0]),
                          int(fragment_indices[i][j][1]),
                          int(fragment_indices[i][j][2]),
                      )
                      if candidate_id == 0:  # this is padding
                          continue
                      prob = round(float(candidate_probs[i][j]), 5)
                      if prob < 0.01:
                          continue
                      # -1 because in the output file we will not have a [CLS] token
                      possible_replacements.append(
                          str(start - 1)
                          + " "
                          + str(end - 1)
                          + " "
                          + str(candidate_id)
                          + " "
                          + str(prob)
                      )
                  all_possible_replacements.append(
                      possible_replacements
                  )
  
              # torch.argmax(tag_logits, dim=-1) gives a tensor of best predicted labels with shape [batch_size, char_seq_len], .dtype = int64
              # character_preds is list of lists of predicted labels
              character_preds = tensor2list(
                  torch.argmax(tag_logits, dim=-1)
              )
              all_tag_preds.extend(character_preds)
  
          if len(all_possible_replacements) != len(
              all_tag_preds
          ) or len(all_possible_replacements) != len(
              infer_datalayer.dataset.examples
          ):
              raise IndexError(
                  "number of sentences mismatch: len(all_possible_replacements)="
                  + str(len(all_possible_replacements))
                  + "; len(all_tag_preds)="
                  + str(len(all_tag_preds))
                  + "; len(infer_datalayer.dataset.examples)="
                  + str(len(infer_datalayer.dataset.examples))
              )
          # save results to file
          with open(output_name, "w", encoding="utf-8") as out:
              for i in range(len(infer_datalayer.dataset.examples)):
                  hyp, ref = infer_datalayer.dataset.hyps_refs[i]
                  num_letters = hyp.count(" ") + 1
                  tag_pred_str = " ".join(
                      list(
                          map(
                              str,
                              all_tag_preds[i][1 : (num_letters + 1)],
                          )
                      )
                  )
                  possible_replacements_str = ";".join(
                      all_possible_replacements[i]
                  )
                  out.write(
                      hyp
                      + "\t"
                      + ref
                      + "\t"
                      + possible_replacements_str
                      + "\t"
                      + tag_pred_str
                      + "\n"
                  )
  
      except Exception as e:
          raise ValueError("Error processing file " + input_name)
  
      finally:
          # set mode back to its original value
          self.train(mode=mode)
          logging.set_verbosity(logging_level)
  
imports: from time import perf_counter
  from typing import Dict, Optional
  import torch
  from omegaconf import DictConfig
  from pytorch_lightning import Trainer
  from nemo.collections.common.losses import CrossEntropyLoss
  from nemo.collections.nlp.data.spellchecking_asr_customization import SpellcheckingAsrCustomizationDataset, SpellcheckingAsrCustomizationTestDataset, TarredSpellcheckingAsrCustomizationDataset, bert_example
  from nemo.collections.nlp.data.text_normalization_as_tagging.utils import read_label_map
  from nemo.collections.nlp.metrics.classification_report import ClassificationReport
  from nemo.collections.nlp.models.nlp_model import NLPModel
  from nemo.collections.nlp.modules.common.token_classifier import TokenClassifier
  from nemo.collections.nlp.parts.utils_funcs import tensor2list
  from nemo.core.classes.common import PretrainedModelInfo, typecheck
  from nemo.core.neural_types import LogitsType, NeuralType
  from nemo.utils import logging
  from nemo.utils.decorators import experimental
label_desc: Use image node_id 11 for calling the SpellcheckingAsrCustomizationModel obj's underlying member method code with example usage: obj.infer(dataloader_cfg, input_name, output_name) without return types