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: Sets up training data: creates dataset and sets data loader. If parameter ``train_data_config`` is not
  provided, then :ref:`config<model-config-label>` section ``train_ds`` will be used.
  
  Args:
      train_data_config (:obj:`Union[Dict[str, Any], DictConfig]`, `optional`): a dictionary that should contain
          only fields present in :ref:`data config<data-config-label>`.
          If some of the fields are missing, then they will be set according to
          :ref:`data config<data-config-label>` defaults. If ``train_data_config`` parameter is not set, then
          ``train_ds`` item of model config is used. Here model config is a configuration used for model
          instantiation.
docstr_tok: ["Sets","up","training","data",":","creates","dataset","and","sets","data","loader",".","If","parameter","``","train_data_config","``","is","not","provided",",","then",":","ref",":","`","config","<","model-config-label",">","`","section","``","train_ds","``","will","be","used",".","Args",":","train_data_config","(",":","obj",":","`","Union","[","Dict","[","str",",","Any","]",",","DictConfig","]","`",",","`","optional","`",")",":","a","dictionary","that","should","contain","only","fields","present","in",":","ref",":","`","data","config","<","data-config-label",">","`",".","If","some","of","the","fields","are","missing",",","then","they","will","be","set","according","to",":","ref",":","`","data","config","<","data-config-label",">","`","defaults",".","If","``","train_data_config","``","parameter","is","not","set",",","then","``","train_ds","``","item","of","model","config","is","used",".","Here","model","config","is","a","configuration","used","for","model","instantiation","."]
code: def setup_training_data(
      self,
      train_data_config: Optional[
          Union[Dict[str, Any], DictConfig]
      ] = None,
  ) -> None:
      """
      Sets up training data: creates dataset and sets data loader. If parameter ``train_data_config`` is not
      provided, then :ref:`config<model-config-label>` section ``train_ds`` will be used.
  
      Args:
          train_data_config (:obj:`Union[Dict[str, Any], DictConfig]`, `optional`): a dictionary that should contain
              only fields present in :ref:`data config<data-config-label>`.
              If some of the fields are missing, then they will be set according to
              :ref:`data config<data-config-label>` defaults. If ``train_data_config`` parameter is not set, then
              ``train_ds`` item of model config is used. Here model config is a configuration used for model
              instantiation.
      """
      if train_data_config is not None:
          train_data_config = OmegaConf.create(train_data_config)
          train_data_config = OmegaConf.merge(
              OmegaConf.structured(
                  PunctuationCapitalizationTrainDataConfig
              ),
              train_data_config,
          )
      if train_data_config is None:
          train_data_config = self._cfg.train_ds
  
      self._train_dl = self._setup_dataloader_from_config(
          cfg=train_data_config, train=True
      )
  
      # Need to set this because if using an IterableDataset, the length of the dataloader is the total number
      # of samples rather than the number of batches, and this messes up the tqdm progress bar.
      # So we set the number of steps manually (to the correct number) to fix this.
      if (
          "use_tarred_dataset" in train_data_config
          and train_data_config["use_tarred_dataset"]
      ):
          # We also need to check if limit_train_batches is already set.
          # If it's an int, we assume that the user has set it to something sane, i.e. <= # training batches,
          # and don't change it. Otherwise, adjust batches accordingly if it's a float (including 1.0).
          if self._trainer is not None and isinstance(
              self._trainer.limit_train_batches, float
          ):
              self._trainer.limit_train_batches = int(
                  self._trainer.limit_train_batches
                  * ceil(len(self._train_dl.dataset) / self.world_size)
              )
          elif self._trainer is None:
              logging.warning(
                  "Model Trainer was not set before constructing the dataset, incorrect number of "
                  "training batches will be used. Please set the trainer and rebuild the dataset."
              )
  
      self.punct_label_ids = (
          self._train_dl.dataset.punct_label_ids.copy()
      )
      self.capit_label_ids = (
          self._train_dl.dataset.capit_label_ids.copy()
      )
      self.label_ids_are_set = True
      if (
          not torch.distributed.is_initialized()
          or torch.distributed.get_rank() == 0
      ):
          label_vocab_dir = (
              self._cfg.common_dataset_parameters.label_vocab_dir
          )
          if label_vocab_dir is None:
              (
                  punct_label_ids_file,
                  capit_label_ids_file,
              ) = self._train_dl.dataset.save_labels_and_get_file_paths(
                  self._cfg.class_labels.punct_labels_file,
                  self._cfg.class_labels.capit_labels_file,
              )
          else:
              punct_label_ids_file = (
                  Path(label_vocab_dir).expanduser()
                  / self._cfg.class_labels.punct_labels_file
              )
              capit_label_ids_file = (
                  Path(label_vocab_dir).expanduser()
                  / self._cfg.class_labels.capit_labels_file
              )
          self.register_artifact(
              "class_labels.punct_labels_file",
              str(punct_label_ids_file),
          )
          self.register_artifact(
              "class_labels.capit_labels_file",
              str(capit_label_ids_file),
          )
  
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 14 for calling the PunctuationCapitalizationModel obj's underlying member method code with example usage: obj.setup_training_data(train_data_config) without return types