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: Same BERT-based model is used to calculate embeddings for sequence of single characters and for sequence of subwords.
  Then we concatenate subword embeddings to each character corresponding to this subword.
  We return logits for each character x 11 labels: 0 - character doesn't belong to any candidate, 1..10 - character belongs to candidate with this id.
  
  # Arguments
      input_ids: token_ids for single characters; .shape = [batch_size, char_seq_len]; .dtype = int64
      input_mask: mask for input_ids(1 - real, 0 - padding); .shape = [batch_size, char_seq_len]; .dtype = int64
      segment_ids: segment types for input_ids (0 - ASR-hypothesis, 1..10 - candidate); .shape = [batch_size, char_seq_len]; .dtype = int64
      input_ids_for_subwords: token_ids for subwords; .shape = [batch_size, subword_seq_len]; .dtype = int64
      input_mask_for_subwords: mask for input_ids_for_subwords(1 - real, 0 - padding); .shape = [batch_size, subword_seq_len]; .dtype = int64
      segment_ids_for_subwords: segment types for input_ids_for_subwords (0 - ASR-hypothesis, 1..10 - candidate); .shape = [batch_size, subword_seq_len]; .dtype = int64
      character_pos_to_subword_pos: tensor mapping character position in the input sequence to subword position; .shape = [batch_size, char_seq_len]; .dtype = int64
docstr_tok: ["Same","BERT-based","model","is","used","to","calculate","embeddings","for","sequence","of","single","characters","and","for","sequence","of","subwords",".","Then","we","concatenate","subword","embeddings","to","each","character","corresponding","to","this","subword",".","We","return","logits","for","each","character","x","11","labels",":","0","-","character","does","n't","belong","to","any","candidate",",","1","..","10","-","character","belongs","to","candidate","with","this","id",".","#","Arguments","input_ids",":","token_ids","for","single","characters",";",".shape","=","[","batch_size",",","char_seq_len","]",";",".dtype","=","int64","input_mask",":","mask","for","input_ids","(","1","-","real",",","0","-","padding",")",";",".shape","=","[","batch_size",",","char_seq_len","]",";",".dtype","=","int64","segment_ids",":","segment","types","for","input_ids","(","0","-","ASR-hypothesis",",","1","..","10","-","candidate",")",";",".shape","=","[","batch_size",",","char_seq_len","]",";",".dtype","=","int64","input_ids_for_subwords",":","token_ids","for","subwords",";",".shape","=","[","batch_size",",","subword_seq_len","]",";",".dtype","=","int64","input_mask_for_subwords",":","mask","for","input_ids_for_subwords","(","1","-","real",",","0","-","padding",")",";",".shape","=","[","batch_size",",","subword_seq_len","]",";",".dtype","=","int64","segment_ids_for_subwords",":","segment","types","for","input_ids_for_subwords","(","0","-","ASR-hypothesis",",","1","..","10","-","candidate",")",";",".shape","=","[","batch_size",",","subword_seq_len","]",";",".dtype","=","int64","character_pos_to_subword_pos",":","tensor","mapping","character","position","in","the","input","sequence","to","subword","position",";",".shape","=","[","batch_size",",","char_seq_len","]",";",".dtype","=","int64"]
code: def forward(
      self,
      input_ids,
      input_mask,
      segment_ids,
      input_ids_for_subwords,
      input_mask_for_subwords,
      segment_ids_for_subwords,
      character_pos_to_subword_pos,
  ):
      """
      Same BERT-based model is used to calculate embeddings for sequence of single characters and for sequence of subwords.
      Then we concatenate subword embeddings to each character corresponding to this subword.
      We return logits for each character x 11 labels: 0 - character doesn't belong to any candidate, 1..10 - character belongs to candidate with this id.
  
      # Arguments
          input_ids: token_ids for single characters; .shape = [batch_size, char_seq_len]; .dtype = int64
          input_mask: mask for input_ids(1 - real, 0 - padding); .shape = [batch_size, char_seq_len]; .dtype = int64
          segment_ids: segment types for input_ids (0 - ASR-hypothesis, 1..10 - candidate); .shape = [batch_size, char_seq_len]; .dtype = int64
          input_ids_for_subwords: token_ids for subwords; .shape = [batch_size, subword_seq_len]; .dtype = int64
          input_mask_for_subwords: mask for input_ids_for_subwords(1 - real, 0 - padding); .shape = [batch_size, subword_seq_len]; .dtype = int64
          segment_ids_for_subwords: segment types for input_ids_for_subwords (0 - ASR-hypothesis, 1..10 - candidate); .shape = [batch_size, subword_seq_len]; .dtype = int64
          character_pos_to_subword_pos: tensor mapping character position in the input sequence to subword position; .shape = [batch_size, char_seq_len]; .dtype = int64
      """
  
      # src_hiddens.shape = [batch_size, char_seq_len, bert_hidden_size]; .dtype=float32
      src_hiddens = self.bert_model(
          input_ids=input_ids,
          token_type_ids=segment_ids,
          attention_mask=input_mask,
      )
      # src_hiddens_for_subwords.shape = [batch_size, subword_seq_len, bert_hidden_size]; .dtype=float32
      src_hiddens_for_subwords = self.bert_model(
          input_ids=input_ids_for_subwords,
          token_type_ids=segment_ids_for_subwords,
          attention_mask=input_mask_for_subwords,
      )
  
      # Next three commands concatenate subword embeddings to each character embedding of the corresponding subword
      # index.shape = [batch_size, char_seq_len, bert_hidden_size]; .dtype=int64
      index = character_pos_to_subword_pos.unsqueeze(-1).expand(
          (-1, -1, src_hiddens_for_subwords.shape[2])
      )
      # src_hiddens_2.shape = [batch_size, char_seq_len, bert_hidden_size]; .dtype=float32
      src_hiddens_2 = torch.gather(src_hiddens_for_subwords, 1, index)
      # src_hiddens.shape = [batch_size, char_seq_len, bert_hidden_size * 2]; .dtype=float32
      src_hiddens = torch.cat((src_hiddens, src_hiddens_2), 2)
  
      # logits.shape = [batch_size, char_seq_len, num_labels]; num_labels=11: ids from 0 to 10; .dtype=float32
      logits = self.logits(hidden_states=src_hiddens)
      return logits
  
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 5 for calling the SpellcheckingAsrCustomizationModel obj's underlying member method code with example usage: obj.forward(input_ids, input_mask, segment_ids, input_ids_for_subwords, input_mask_for_subwords, segment_ids_for_subwords, character_pos_to_subword_pos) and returns: logits