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: Creates dataset and data loader according to config ``cfg``. If ``train=False`` and attributes
  ``punct_label_ids`` and ``capit_label_ids`` are not set, then this method sets the attributes and registers
  label artifacts.
  
  Args:
      cfg (:obj:`DictConfig`): a config which follows dataclass
          :class:`~nemo.collections.nlp.data.token_classification.punctuation_capitalization_dataset.PunctuationCapitalizationEvalDataConfig`
          Note that list ``ds_item`` is not supported because list ``ds_item`` is unpacked by NeMo core
          instruments
      train (:obj:`bool`): whether train data is set. If ``True``, then label ids are not set in this function
docstr_tok: ["Creates","dataset","and","data","loader","according","to","config","``","cfg","``",".","If","``","train=False","``","and","attributes","``","punct_label_ids","``","and","``","capit_label_ids","``","are","not","set",",","then","this","method","sets","the","attributes","and","registers","label","artifacts",".","Args",":","cfg","(",":","obj",":","`","DictConfig","`",")",":","a","config","which","follows","dataclass",":","class",":","`","~nemo.collections.nlp.data.token_classification.punctuation_capitalization_dataset.PunctuationCapitalizationEvalDataConfig","`","Note","that","list","``","ds_item","``","is","not","supported","because","list","``","ds_item","``","is","unpacked","by","NeMo","core","instruments","train","(",":","obj",":","`","bool","`",")",":","whether","train","data","is","set",".","If","``","True","``",",","then","label","ids","are","not","set","in","this","function"]
code: def _setup_dataloader_from_config(
      self, cfg: DictConfig, train: bool
  ) -> torch.utils.data.DataLoader:
      """
      Creates dataset and data loader according to config ``cfg``. If ``train=False`` and attributes
      ``punct_label_ids`` and ``capit_label_ids`` are not set, then this method sets the attributes and registers
      label artifacts.
  
      Args:
          cfg (:obj:`DictConfig`): a config which follows dataclass
              :class:`~nemo.collections.nlp.data.token_classification.punctuation_capitalization_dataset.PunctuationCapitalizationEvalDataConfig`
              Note that list ``ds_item`` is not supported because list ``ds_item`` is unpacked by NeMo core
              instruments
          train (:obj:`bool`): whether train data is set. If ``True``, then label ids are not set in this function
      """
      self._check_label_config_parameters()
      if not self.label_ids_are_set and not train:
          self._set_label_ids()
      if cfg.use_tarred_dataset:
          if cfg.tar_metadata_file is None:
              raise ValueError(
                  f"If parameter `use_tarred_dataset` is `True`, then a field `tar_metadata_file` has to be a path "
                  f"to tarred dataset metadata file, whereas `None` is given."
              )
          tar_metadata_file = Path(cfg.ds_item) / cfg.tar_metadata_file
          dataset = BertPunctuationCapitalizationTarredDataset(
              metadata_file=tar_metadata_file,
              tokenizer=self.tokenizer,
              pad_label=self._cfg.common_dataset_parameters.pad_label,
              ignore_extra_tokens=self._cfg.common_dataset_parameters.ignore_extra_tokens,
              ignore_start_end=self._cfg.common_dataset_parameters.ignore_start_end,
              world_size=self.world_size,
              global_rank=self.global_rank,
              shuffle_n=cfg.tar_shuffle_n,
              shard_strategy=cfg.shard_strategy,
              label_info_save_dir=cfg.label_info_save_dir,
              use_audio=cfg.use_audio,
          )
          dataset.check_for_label_consistency_with_model_config(
              self.punct_label_ids,
              self.capit_label_ids,
              self._cfg.class_labels,
              self._cfg.common_dataset_parameters,
          )
      else:
          if cfg.text_file is None or cfg.labels_file is None:
              raise ValueError(
                  f"If parameter `use_tarred_dataset` is `False`, then fields `text_file` and `labels_file` in "
                  f"dataset config must not be `None`. Whereas `text_file={cfg.text_file}` and "
                  f"`label_file={cfg.labels_file}`."
              )
          if cfg.tokens_in_batch is None and cfg.use_bucketing:
              raise ValueError(
                  f"If `use_tarred_dataset` is `False`, then you need to provide `tokens_in_batch` parameter."
              )
          (
              text_file,
              labels_file,
          ) = (
              Path(cfg.ds_item) / cfg.text_file,
              Path(cfg.ds_item) / cfg.labels_file,
          )
          if cfg.audio_file:
              audio_file = Path(cfg.ds_item) / cfg.audio_file
          if self.label_ids_are_set:
              label_kwargs = {
                  "punct_label_ids": self.punct_label_ids,
                  "capit_label_ids": self.capit_label_ids,
              }
          else:
              (
                  punct_label_vocab_file,
                  capit_label_vocab_file,
              ) = self._extract_label_vocab_files_from_config()
              label_kwargs = {
                  "punct_label_ids": self._cfg.common_dataset_parameters.punct_label_ids,
                  "capit_label_ids": self._cfg.common_dataset_parameters.capit_label_ids,
                  "punct_label_vocab_file": punct_label_vocab_file,
                  "capit_label_vocab_file": capit_label_vocab_file,
              }
          if train:
              number_of_batches_is_multiple_of = 1
              if self._trainer is None:
                  warnings.warn(
                      "A model attribute `trainer` is not set before training dataset setting. If training is "
                      "resumed from checkpoint, then current epoch data loading can be distorted: some batches "
                      "may be processed several times and some can be not processed at all. `trainer.current_epoch`"
                      " is used as random seed for shuffling batches. Now 0 will be used. If the "
                      "checkpoint was created not during initial epoch a shuffling of the dataset will "
                      "be different. You may try use `exp_manager()` function and "
                      "`PunctuationCapitalizationModel.set_trainer()` method before "
                      "`PunctuationCapitalizationModel.setup_training_data()` method."
                  )
                  batch_shuffling_random_seed = 0
              else:
                  batch_shuffling_random_seed = (
                      self._trainer.current_epoch
                  )
          else:
              batch_shuffling_random_seed = 0
              if self._trainer is None:
                  warnings.warn(
                      "A model attribute `trainer` is not set before test or validation dataset setting. If more "
                      "than 1 GPU is used for testing, then some examples may be tested several times because "
                      "number of batches may be not evenly divisible by number of processes. This leads to "
                      "distortion of metrics. See more in description of `number_of_batches_is_multiple_of` "
                      "parameter of class `BertPunctuationCapitalizationDataset` initializer and "
                      "https://pytorch.org/docs/stable/data.html#multi-process-data-loading. You may try to use "
                      "`PunctuationCapitalizationModel.set_trainer()` method before "
                      "`PunctuationCapitalizationModel.setup_validation_data()` and "
                      "`PunctuationCapitalizationModel.setup_test_data()` methods."
                  )
                  number_of_batches_is_multiple_of = 1
              else:
                  number_of_batches_is_multiple_of = (
                      self._trainer.num_nodes
                      * self._trainer.num_devices
                  )
          if cfg.cache_dir is None:
              cache_dir = cfg.cache_dir
          else:
              # If pickled features are saved `cache_dir` not in the same directory with original data files, then
              # a full path to data directory have to be appended to `cache_dir`. This is done to avoid collisions
              # cache for different datasets is saved to same `cache_dir`.
              cache_dir = Path(cfg.cache_dir).joinpath(
                  "fsroot",
                  *text_file.expanduser().resolve().parts[1:-1],
              )
          dataset = BertPunctuationCapitalizationDataset(
              tokenizer=self.tokenizer,
              text_file=text_file,
              labels_file=labels_file,
              pad_label=self._cfg.common_dataset_parameters.pad_label,
              **label_kwargs,
              max_seq_length=cfg.max_seq_length,
              ignore_extra_tokens=self._cfg.common_dataset_parameters.ignore_extra_tokens,
              ignore_start_end=self._cfg.common_dataset_parameters.ignore_start_end,
              use_cache=cfg.use_cache,
              num_samples=cfg.num_samples,
              tokens_in_batch=cfg.tokens_in_batch,
              n_jobs=cfg.n_jobs,
              number_of_batches_is_multiple_of=number_of_batches_is_multiple_of,
              batch_shuffling_random_seed=batch_shuffling_random_seed,
              verbose=cfg.verbose,
              get_label_frequencies=cfg.get_label_frequences,
              cache_dir=cache_dir,
              label_info_save_dir=cfg.label_info_save_dir,
              audio_file=audio_file if cfg.audio_file else None,
              sample_rate=cfg.sample_rate,
              use_audio=cfg.use_audio,
              use_bucketing=cfg.use_bucketing,
              preload_audios=cfg.preload_audios,
          )
      if cfg.shuffle and cfg.use_tarred_dataset:
          logging.warning(
              f"Shuffling in dataloader is not supported for tarred dataset."
          )
          shuffle = False
      else:
          shuffle = cfg.shuffle
      return torch.utils.data.DataLoader(
          dataset=dataset,
          collate_fn=dataset.collate_fn,
          batch_size=1 if cfg.use_bucketing else cfg.batch_size,
          shuffle=shuffle,
          num_workers=cfg.num_workers,
          pin_memory=cfg.pin_memory,
          drop_last=cfg.drop_last,
          persistent_workers=cfg.persistent_workers
          if cfg.num_workers > 0
          else False,
      )
  
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 22 for calling the PunctuationCapitalizationModel obj's underlying member method code with example usage: obj._setup_dataloader_from_config(cfg, train) and returns: torch