class_docstr: A model for restoring punctuation and capitalization in text using lexical and audio features.
  
  The model consists of a language model and two multilayer perceptrons (MLP) on top the fusion of LM and AM. The first
  MLP serves for punctuation prediction and the second is for capitalization prediction. You can use only BERT-like
  HuggingFace language models (model ``forward`` method accepts ``input_ids``, ``token_types_ids``,
  ``attention_mask`` arguments). See more about model config options :ref:`here<model-config-label>`.
  And any :class:`~nemo.collections.asr.models.EncDecCTCModel` which has encoder module which is used as an AM.
  
  For training and testing use dataset
  :class:`~nemo.collections.nlp.data.token_classification.punctuation_capitalization_dataset.BertPunctuationCapitalizationDataset` with parameter ``use_audio`` set to ``True``,
  for training on huge amounts of data which cannot be loaded into memory simultaneously use
  :class:`~nemo.collections.nlp.data.token_classification.punctuation_capitalization_tarred_dataset.BertPunctuationCapitalizationTarredDataset` with parameter ``use_audio`` set to ``True``.
  
  Args:
      cfg: a model configuration. It should follow dataclass
          :class:`~nemo.collections.nlp.models.token_classification.punctuation_capitalization_config.PunctuationCapitalizationLexicalAudioModelConfig`
          See an example of full config in
          `nemo/examples/nlp/token_classification/conf/punctuation_capitalization_lexical_audio_config.yaml
          <https://github.com/NVIDIA/NeMo/blob/main/examples/nlp/token_classification/conf/punctuation_capitalization_lexical_audio_config.yaml>`_
      trainer: an instance of a PyTorch Lightning trainer
class_docstr_tok: ["A","model","for","restoring","punctuation","and","capitalization","in","text","using","lexical","and","audio","features",".","The","model","consists","of","a","language","model","and","two","multilayer","perceptrons","(","MLP",")","on","top","the","fusion","of","LM","and","AM",".","The","first","MLP","serves","for","punctuation","prediction","and","the","second","is","for","capitalization","prediction",".","You","can","use","only","BERT-like","HuggingFace","language","models","(","model","``","forward","``","method","accepts","``","input_ids","``",",","``","token_types_ids","``",",","``","attention_mask","``","arguments",")",".","See","more","about","model","config","options",":","ref",":","`","here","<","model-config-label",">","`",".","And","any",":","class",":","`","~nemo.collections.asr.models.EncDecCTCModel","`","which","has","encoder","module","which","is","used","as","an","AM",".","For","training","and","testing","use","dataset",":","class",":","`","~nemo.collections.nlp.data.token_classification.punctuation_capitalization_dataset.BertPunctuationCapitalizationDataset","`","with","parameter","``","use_audio","``","set","to","``","True","``",",","for","training","on","huge","amounts","of","data","which","can","not","be","loaded","into","memory","simultaneously","use",":","class",":","`","~nemo.collections.nlp.data.token_classification.punctuation_capitalization_tarred_dataset.BertPunctuationCapitalizationTarredDataset","`","with","parameter","``","use_audio","``","set","to","``","True","``",".","Args",":","cfg",":","a","model","configuration",".","It","should","follow","dataclass",":","class",":","`","~nemo.collections.nlp.models.token_classification.punctuation_capitalization_config.PunctuationCapitalizationLexicalAudioModelConfig","`","See","an","example","of","full","config","in","`","nemo\/examples\/nlp\/token_classification\/conf\/punctuation_capitalization_lexical_audio_config.yaml","<","https",":","\/\/github.com\/NVIDIA\/NeMo\/blob\/main\/examples\/nlp\/token_classification\/conf\/punctuation_capitalization_lexical_audio_config.yaml",">","`","_","trainer",":","an","instance","of","a","PyTorch","Lightning","trainer"]
docstr: Executes a forward pass through the model. For more details see ``forward`` method of :class:`~nemo.collections.nlp.models.token_classification.punctuation_capitalization_config.PunctuationCapitalizationLexicalAudioModelConfig`
  and ``forward`` method of :class:'~nemo.collections.asr.models.EncDecCTCModel'
  
  Args:
      input_ids (:obj:`torch.Tensor`): an integer torch tensor of shape ``[Batch, Time]``. Contains encoded
          source tokens.
      attention_mask (:obj:`torch.Tensor`): a boolean torch tensor of shape ``[Batch, Time]``. Contains an
          attention mask for excluding paddings.
      token_type_ids (:obj:`torch.Tensor`): an integer torch Tensor of shape ``[Batch, Time]``. Contains an index
          of segment to which a token belongs. If ``token_type_ids`` is not ``None``, then it should be a zeros
          tensor.
      features (:obj:`torch.Tensor`): tensor that represents a batch of raw audio signals,
          of shape [B, T]. T here represents timesteps, with 1 second of audio represented as
          sample_rate number of floating point values.
      features_length (:obj:`torch.Tensor`): Vector of length B, that contains the individual lengths of the audio
          sequences.
  
  Returns:
      :obj:`Tuple[torch.Tensor, torch.Tensor]`: a tuple containing
  
          - ``punct_logits`` (:obj:`torch.Tensor`): a float torch tensor of shape
            ``[Batch, Time, NumPunctuationLabels]`` containing punctuation logits
          - ``capit_logits`` (:obj:`torch.Tensor`): a float torch tensor of shape
            ``[Batch, Time, NumCapitalizationLabels]`` containing capitalization logits
docstr_tok: ["Executes","a","forward","pass","through","the","model",".","For","more","details","see","``","forward","``","method","of",":","class",":","`","~nemo.collections.nlp.models.token_classification.punctuation_capitalization_config.PunctuationCapitalizationLexicalAudioModelConfig","`","and","``","forward","``","method","of",":","class",":","'~nemo.collections.asr.models.EncDecCTCModel'","Args",":","input_ids","(",":","obj",":","`","torch.Tensor","`",")",":","an","integer","torch","tensor","of","shape","``","[","Batch",",","Time","]","``",".","Contains","encoded","source","tokens",".","attention_mask","(",":","obj",":","`","torch.Tensor","`",")",":","a","boolean","torch","tensor","of","shape","``","[","Batch",",","Time","]","``",".","Contains","an","attention","mask","for","excluding","paddings",".","token_type_ids","(",":","obj",":","`","torch.Tensor","`",")",":","an","integer","torch","Tensor","of","shape","``","[","Batch",",","Time","]","``",".","Contains","an","index","of","segment","to","which","a","token","belongs",".","If","``","token_type_ids","``","is","not","``","None","``",",","then","it","should","be","a","zeros","tensor",".","features","(",":","obj",":","`","torch.Tensor","`",")",":","tensor","that","represents","a","batch","of","raw","audio","signals",",","of","shape","[","B",",","T","]",".","T","here","represents","timesteps",",","with","1","second","of","audio","represented","as","sample_rate","number","of","floating","point","values",".","features_length","(",":","obj",":","`","torch.Tensor","`",")",":","Vector","of","length","B",",","that","contains","the","individual","lengths","of","the","audio","sequences",".","Returns",":",":","obj",":","`","Tuple","[","torch.Tensor",",","torch.Tensor","]","`",":","a","tuple","containing","-","``","punct_logits","``","(",":","obj",":","`","torch.Tensor","`",")",":","a","float","torch","tensor","of","shape","``","[","Batch",",","Time",",","NumPunctuationLabels","]","``","containing","punctuation","logits","-","``","capit_logits","``","(",":","obj",":","`","torch.Tensor","`",")",":","a","float","torch","tensor","of","shape","``","[","Batch",",","Time",",","NumCapitalizationLabels","]","``","containing","capitalization","logits"]
code: def forward(
      self,
      input_ids: torch.Tensor,
      attention_mask: torch.Tensor,
      token_type_ids: Optional[torch.Tensor] = None,
      features: torch.Tensor = None,
      features_length: torch.Tensor = None,
  ) -> Tuple[torch.Tensor, torch.Tensor]:
      """
      Executes a forward pass through the model. For more details see ``forward`` method of :class:`~nemo.collections.nlp.models.token_classification.punctuation_capitalization_config.PunctuationCapitalizationLexicalAudioModelConfig`
      and ``forward`` method of :class:'~nemo.collections.asr.models.EncDecCTCModel'
  
      Args:
          input_ids (:obj:`torch.Tensor`): an integer torch tensor of shape ``[Batch, Time]``. Contains encoded
              source tokens.
          attention_mask (:obj:`torch.Tensor`): a boolean torch tensor of shape ``[Batch, Time]``. Contains an
              attention mask for excluding paddings.
          token_type_ids (:obj:`torch.Tensor`): an integer torch Tensor of shape ``[Batch, Time]``. Contains an index
              of segment to which a token belongs. If ``token_type_ids`` is not ``None``, then it should be a zeros
              tensor.
          features (:obj:`torch.Tensor`): tensor that represents a batch of raw audio signals,
              of shape [B, T]. T here represents timesteps, with 1 second of audio represented as
              sample_rate number of floating point values.
          features_length (:obj:`torch.Tensor`): Vector of length B, that contains the individual lengths of the audio
              sequences.
  
      Returns:
          :obj:`Tuple[torch.Tensor, torch.Tensor]`: a tuple containing
  
              - ``punct_logits`` (:obj:`torch.Tensor`): a float torch tensor of shape
                ``[Batch, Time, NumPunctuationLabels]`` containing punctuation logits
              - ``capit_logits`` (:obj:`torch.Tensor`): a float torch tensor of shape
                ``[Batch, Time, NumCapitalizationLabels]`` containing capitalization logits
      """
      self.update_max_seq_length(
          seq_length=features.size(1), device=features.device
      )
      lexical_hidden_states = self.bert_model(
          input_ids=input_ids,
          token_type_ids=token_type_ids,
          attention_mask=attention_mask,
      )
      if isinstance(lexical_hidden_states, tuple):
          lexical_hidden_states = lexical_hidden_states[0]
  
      (
          processed_signal,
          processed_signal_length,
      ) = self.audio_encoder.preprocessor(
          input_signal=features,
          length=features_length,
      )
  
      if (
          self.audio_encoder.spec_augmentation is not None
          and self.training
      ):
          processed_signal = self.audio_encoder.spec_augmentation(
              input_spec=processed_signal,
              length=processed_signal_length,
          )
  
      (
          audio_hidden_states,
          audio_hidden_states_length,
      ) = self.audio_encoder.encoder(
          audio_signal=processed_signal, length=processed_signal_length
      )
      audio_hidden_states = audio_hidden_states.permute(0, 2, 1)
      audio_hidden_states = self.audio_proj(audio_hidden_states)
  
      fused = self.fusion(
          lexical_hidden_states,
          attention_mask,
          audio_hidden_states,
          self.make_pad_mask(
              audio_hidden_states.size(1), audio_hidden_states_length
          ),
      )
  
      punct_logits = self.punct_classifier(hidden_states=fused)
      capit_logits = self.capit_classifier(hidden_states=fused)
  
      return punct_logits, capit_logits
  
imports: import os
  from math import ceil
  from typing import Any, Dict, List, Optional, Tuple, Union
  import numpy
  import torch
  from omegaconf import DictConfig, open_dict
  from pytorch_lightning import Trainer
  from torch.nn import Linear
  from tqdm import tqdm
  from nemo.collections.common.losses.cross_entropy import CrossEntropyLoss
  from nemo.collections.nlp.models.token_classification.punctuation_capitalization_model import PunctuationCapitalizationModel
  from nemo.collections.nlp.modules.common.transformer import TransformerDecoder
  from nemo.core.classes.common import PretrainedModelInfo
  from nemo.core.classes.mixins import adapter_mixins
  from nemo.utils import logging
label_desc: Use image node_id 3 for calling the PunctuationCapitalizationLexicalAudioModel obj's underlying member method code with example usage: obj.forward(input_ids, attention_mask, token_type_ids, features, features_length) and returns: punct_logits,  capit_logits