class_docstr: A model for restoring punctuation and capitalization in text. The model is usually used together with ASR model
  because ASR models often return text without punctuation and capitalization.
  
  The model consists of a language model and two multilayer perceptrons (MLP) on top the language model. 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>`.
  
  Use method :meth:`~add_punctuation_capitalization` for model inference.
  
  For training and testing use dataset
  :class:`~nemo.collections.nlp.data.token_classification.punctuation_capitalization_dataset.BertPunctuationCapitalizationDataset`,
  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`.
  
  Args:
      cfg: a model configuration. It should follow dataclass
          :class:`~nemo.collections.nlp.models.token_classification.punctuation_capitalization_config.PunctuationCapitalizationModelConfig`
          See an example of full config in
          `nemo/examples/nlp/token_classification/conf/punctuation_capitalization_config.yaml
          <https://github.com/NVIDIA/NeMo/blob/main/examples/nlp/token_classification/conf/punctuation_capitalization_config.yaml>`_
      trainer: an instance of a PyTorch Lightning trainer
class_docstr_tok: ["A","model","for","restoring","punctuation","and","capitalization","in","text",".","The","model","is","usually","used","together","with","ASR","model","because","ASR","models","often","return","text","without","punctuation","and","capitalization",".","The","model","consists","of","a","language","model","and","two","multilayer","perceptrons","(","MLP",")","on","top","the","language","model",".","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",">","`",".","Use","method",":","meth",":","`","~add_punctuation_capitalization","`","for","model","inference",".","For","training","and","testing","use","dataset",":","class",":","`","~nemo.collections.nlp.data.token_classification.punctuation_capitalization_dataset.BertPunctuationCapitalizationDataset","`",",","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","`",".","Args",":","cfg",":","a","model","configuration",".","It","should","follow","dataclass",":","class",":","`","~nemo.collections.nlp.models.token_classification.punctuation_capitalization_config.PunctuationCapitalizationModelConfig","`","See","an","example","of","full","config","in","`","nemo\/examples\/nlp\/token_classification\/conf\/punctuation_capitalization_config.yaml","<","https",":","\/\/github.com\/NVIDIA\/NeMo\/blob\/main\/examples\/nlp\/token_classification\/conf\/punctuation_capitalization_config.yaml",">","`","_","trainer",":","an","instance","of","a","PyTorch","Lightning","trainer"]
docstr: Applies softmax to get punctuation and capitalization probabilities, applies ``subtokens_mask`` to extract
  probabilities for words from probabilities for tokens, removes ``margin`` probabilities near edges of a segment.
  Left margin of the first segment in a query and right margin of the last segment in a query are not removed.
  Calculates new ``start_word_ids`` taking into the account the margins. If the left margin of a segment is
  removed corresponding start word index is increased by number of words (number of nonzero values in
  corresponding ``subtokens_mask``) in the margin.
  
  Args:
      punct_logits: a float tensor of shape ``[batch_size, segment_length, number_of_punctuation_labels]``
      capit_logits: a float tensor of shape ``[batch_size, segment_length, number_of_capitalization_labels]``
      subtokens_mask: a float tensor of shape ``[batch_size, segment_length]``
      start_word_ids: indices of segment first words in a query
      margin: number of tokens near edges of a segment which probabilities are discarded
      is_first: is segment the first segment in a query
      is_last: is segment the last segment in a query
  Returns:
      b_punct_probs: list containing ``batch_size`` numpy arrays. The numpy arrays have shapes
          ``[number_of_word_in_this_segment, number_of_punctuation_labels]``. Word punctuation probabilities for
          segments in the batch.
      b_capit_probs: list containing ``batch_size`` numpy arrays. The numpy arrays have shapes
          ``[number_of_word_in_this_segment, number_of_capitalization_labels]``. Word capitalization
          probabilities for segments in the batch.
      new_start_word_ids: indices of segment first words in a query after margin removal
docstr_tok: ["Applies","softmax","to","get","punctuation","and","capitalization","probabilities",",","applies","``","subtokens_mask","``","to","extract","probabilities","for","words","from","probabilities","for","tokens",",","removes","``","margin","``","probabilities","near","edges","of","a","segment",".","Left","margin","of","the","first","segment","in","a","query","and","right","margin","of","the","last","segment","in","a","query","are","not","removed",".","Calculates","new","``","start_word_ids","``","taking","into","the","account","the","margins",".","If","the","left","margin","of","a","segment","is","removed","corresponding","start","word","index","is","increased","by","number","of","words","(","number","of","nonzero","values","in","corresponding","``","subtokens_mask","``",")","in","the","margin",".","Args",":","punct_logits",":","a","float","tensor","of","shape","``","[","batch_size",",","segment_length",",","number_of_punctuation_labels","]","``","capit_logits",":","a","float","tensor","of","shape","``","[","batch_size",",","segment_length",",","number_of_capitalization_labels","]","``","subtokens_mask",":","a","float","tensor","of","shape","``","[","batch_size",",","segment_length","]","``","start_word_ids",":","indices","of","segment","first","words","in","a","query","margin",":","number","of","tokens","near","edges","of","a","segment","which","probabilities","are","discarded","is_first",":","is","segment","the","first","segment","in","a","query","is_last",":","is","segment","the","last","segment","in","a","query","Returns",":","b_punct_probs",":","list","containing","``","batch_size","``","numpy","arrays",".","The","numpy","arrays","have","shapes","``","[","number_of_word_in_this_segment",",","number_of_punctuation_labels","]","``",".","Word","punctuation","probabilities","for","segments","in","the","batch",".","b_capit_probs",":","list","containing","``","batch_size","``","numpy","arrays",".","The","numpy","arrays","have","shapes","``","[","number_of_word_in_this_segment",",","number_of_capitalization_labels","]","``",".","Word","capitalization","probabilities","for","segments","in","the","batch",".","new_start_word_ids",":","indices","of","segment","first","words","in","a","query","after","margin","removal"]
code: def _transform_logit_to_prob_and_remove_margins_and_extract_word_probs(
      self,
      punct_logits: torch.Tensor,
      capit_logits: torch.Tensor,
      subtokens_mask: torch.Tensor,
      start_word_ids: Tuple[int],
      margin: int,
      is_first: Tuple[bool],
      is_last: Tuple[bool],
  ) -> Tuple[List[np.ndarray], List[np.ndarray], List[int]]:
      """
      Applies softmax to get punctuation and capitalization probabilities, applies ``subtokens_mask`` to extract
      probabilities for words from probabilities for tokens, removes ``margin`` probabilities near edges of a segment.
      Left margin of the first segment in a query and right margin of the last segment in a query are not removed.
      Calculates new ``start_word_ids`` taking into the account the margins. If the left margin of a segment is
      removed corresponding start word index is increased by number of words (number of nonzero values in
      corresponding ``subtokens_mask``) in the margin.
  
      Args:
          punct_logits: a float tensor of shape ``[batch_size, segment_length, number_of_punctuation_labels]``
          capit_logits: a float tensor of shape ``[batch_size, segment_length, number_of_capitalization_labels]``
          subtokens_mask: a float tensor of shape ``[batch_size, segment_length]``
          start_word_ids: indices of segment first words in a query
          margin: number of tokens near edges of a segment which probabilities are discarded
          is_first: is segment the first segment in a query
          is_last: is segment the last segment in a query
      Returns:
          b_punct_probs: list containing ``batch_size`` numpy arrays. The numpy arrays have shapes
              ``[number_of_word_in_this_segment, number_of_punctuation_labels]``. Word punctuation probabilities for
              segments in the batch.
          b_capit_probs: list containing ``batch_size`` numpy arrays. The numpy arrays have shapes
              ``[number_of_word_in_this_segment, number_of_capitalization_labels]``. Word capitalization
              probabilities for segments in the batch.
          new_start_word_ids: indices of segment first words in a query after margin removal
      """
      new_start_word_ids = list(start_word_ids)
      subtokens_mask = subtokens_mask > 0.5
      b_punct_probs, b_capit_probs = [], []
      for i, (first, last, pl, cl, stm) in enumerate(
          zip(
              is_first,
              is_last,
              punct_logits,
              capit_logits,
              subtokens_mask,
          )
      ):
          if not first:
              new_start_word_ids[i] += torch.count_nonzero(
                  stm[: margin + 1]
              ).numpy()  # + 1 is for [CLS] token
          stm = self._remove_margins(
              stm, margin, keep_left=first, keep_right=last
          )
          for b_probs, logits in [
              (b_punct_probs, pl),
              (b_capit_probs, cl),
          ]:
              p = torch.nn.functional.softmax(
                  self._remove_margins(
                      logits, margin, keep_left=first, keep_right=last
                  )[stm],
                  dim=-1,
              )
              b_probs.append(p.detach().cpu().numpy())
      return b_punct_probs, b_capit_probs, new_start_word_ids
  
imports: import copy
  import warnings
  from math import ceil
  from pathlib import Path
  from typing import Any, Dict, List, Optional, Tuple, Union
  import numpy
  import torch
  from omegaconf import DictConfig, OmegaConf
  from pytorch_lightning import Trainer
  from tqdm import tqdm
  from nemo.collections.common.losses import AggregatorLoss, CrossEntropyLoss
  from nemo.collections.common.metrics import GlobalAverageLossMetric
  from nemo.collections.nlp.data.token_classification.punctuation_capitalization_dataset import BertPunctuationCapitalizationDataset, PunctuationCapitalizationEvalDataConfig, PunctuationCapitalizationTrainDataConfig, load_label_ids, raise_not_equal_labels_error
  from nemo.collections.nlp.data.token_classification.punctuation_capitalization_infer_dataset import BertPunctuationCapitalizationInferDataset
  from nemo.collections.nlp.data.token_classification.punctuation_capitalization_tarred_dataset import BertPunctuationCapitalizationTarredDataset
  from nemo.collections.nlp.metrics.classification_report import ClassificationReport
  from nemo.collections.nlp.models.nlp_model import NLPModel
  from nemo.collections.nlp.models.token_classification.punctuation_capitalization_config import is_legacy_model_config, legacy_model_config_to_new_model_config
  from nemo.collections.nlp.modules.common import TokenClassifier
  from nemo.core.classes.common import PretrainedModelInfo, typecheck
  from nemo.core.classes.exportable import Exportable
  from nemo.core.neural_types import LogitsType, NeuralType
  from nemo.utils import logging
label_desc: Use image node_id 25 for calling the PunctuationCapitalizationModel obj's underlying member method code with example usage: obj._transform_logit_to_prob_and_remove_margins_and_extract_word_probs(punct_logits, capit_logits, subtokens_mask, start_word_ids, margin, is_first, is_last) and returns: b_punct_probs,  b_capit_probs,  new_start_word_ids