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: Set new values for some sections of config. Useful after restoring from checkpoint for fine-tuning
  and testing if config parameters of a restored checkpoint are not suitable.
  
  For ``class_labels``, ``common_dataset_parameters``, ``train_ds``, ``validation_ds``, ``test_ds``, there is
  no need to provide values for all items in an updated config section. If an item is omitted in this method
  parameter, then corresponding item in model config does not change.
  
  If the entire updated section is missing in the model config, then omitted items from this method parameters
  are set according to default values listed
  :ref:`here <run-config-label>`.
  
  .. warning::
      Parameter ``optim`` is processed in a special way. ``optim`` contents are used not for updating of
      model config, but for replacement of entire config section.
  
  If one of parameters ``train_ds``, ``validation_ds``, ``test_ds``, is provided but its value is
  ``None``, then corresponding section is replaced with ``None``.
  
  .. warning::
      You may change values of parameters related to label ids:
  
          - ``common_dataset_parameters.punct_label_ids``,
          - ``common_dataset_parameters.capit_label_ids``,
          - ``common_dataset_parameters.label_vocab_dir``,
          - ``class_labels.punct_labels_file``,
          - ``class_labels.capit_labels_file``,
  
      yet label ids in these parameters must be equal to label ids loaded from checkpoint. Otherwise,
      an error will be raised.
  
  Keyword Args:
      class_labels (:obj:`Union[DictConfig, Dict[str, str]]`): names of label id files used as label
          id dictionaries. See more in :ref:`class labels' config<class-labels-config-label>`.
      common_dataset_parameters (:obj:`Union[DictConfig, Dict[str, Any]]`, `optional`): see more in
          :ref:`common dataset parameters config<common-dataset-parameters-config-label>`.
      train_ds (:obj:`Union[DictConfig, Dict[str, Any]]`, `optional`): configuration of training dataset. See
          possible options in :ref:`data config<data-config-label>`.
      validation_ds (:obj:`Union[DictConfig, Dict[str, Any]]`, `optional`): configuration of validation
          dataset. See possible options in :ref:`data config<data-config-label>`.
      test_ds (:obj:`Union[DictConfig, Dict[str, Any]]`, `optional`): configuration of test dataset. See
          possible options in :ref:`data config<data-config-label>`.
      optim (:obj:`Union[DictConfig, Dict[str, Any]]`, `optional`): optimization configuration. See possible
          options in :ref:`optimization<optimization-label>` and in `primer
          <https://github.com/NVIDIA/NeMo/blob/main/tutorials/00_NeMo_Primer.ipynb>`_ tutorial.
docstr_tok: ["Set","new","values","for","some","sections","of","config",".","Useful","after","restoring","from","checkpoint","for","fine-tuning","and","testing","if","config","parameters","of","a","restored","checkpoint","are","not","suitable",".","For","``","class_labels","``",",","``","common_dataset_parameters","``",",","``","train_ds","``",",","``","validation_ds","``",",","``","test_ds","``",",","there","is","no","need","to","provide","values","for","all","items","in","an","updated","config","section",".","If","an","item","is","omitted","in","this","method","parameter",",","then","corresponding","item","in","model","config","does","not","change",".","If","the","entire","updated","section","is","missing","in","the","model","config",",","then","omitted","items","from","this","method","parameters","are","set","according","to","default","values","listed",":","ref",":","`","here","<","run-config-label",">","`",".","..","warning",":",":","Parameter","``","optim","``","is","processed","in","a","special","way",".","``","optim","``","contents","are","used","not","for","updating","of","model","config",",","but","for","replacement","of","entire","config","section",".","If","one","of","parameters","``","train_ds","``",",","``","validation_ds","``",",","``","test_ds","``",",","is","provided","but","its","value","is","``","None","``",",","then","corresponding","section","is","replaced","with","``","None","``",".","..","warning",":",":","You","may","change","values","of","parameters","related","to","label","ids",":","-","``","common_dataset_parameters.punct_label_ids","``",",","-","``","common_dataset_parameters.capit_label_ids","``",",","-","``","common_dataset_parameters.label_vocab_dir","``",",","-","``","class_labels.punct_labels_file","``",",","-","``","class_labels.capit_labels_file","``",",","yet","label","ids","in","these","parameters","must","be","equal","to","label","ids","loaded","from","checkpoint",".","Otherwise",",","an","error","will","be","raised",".","Keyword","Args",":","class_labels","(",":","obj",":","`","Union","[","DictConfig",",","Dict","[","str",",","str","]","]","`",")",":","names","of","label","id","files","used","as","label","id","dictionaries",".","See","more","in",":","ref",":","`","class","labels","'","config","<","class-labels-config-label",">","`",".","common_dataset_parameters","(",":","obj",":","`","Union","[","DictConfig",",","Dict","[","str",",","Any","]","]","`",",","`","optional","`",")",":","see","more","in",":","ref",":","`","common","dataset","parameters","config","<","common-dataset-parameters-config-label",">","`",".","train_ds","(",":","obj",":","`","Union","[","DictConfig",",","Dict","[","str",",","Any","]","]","`",",","`","optional","`",")",":","configuration","of","training","dataset",".","See","possible","options","in",":","ref",":","`","data","config","<","data-config-label",">","`",".","validation_ds","(",":","obj",":","`","Union","[","DictConfig",",","Dict","[","str",",","Any","]","]","`",",","`","optional","`",")",":","configuration","of","validation","dataset",".","See","possible","options","in",":","ref",":","`","data","config","<","data-config-label",">","`",".","test_ds","(",":","obj",":","`","Union","[","DictConfig",",","Dict","[","str",",","Any","]","]","`",",","`","optional","`",")",":","configuration","of","test","dataset",".","See","possible","options","in",":","ref",":","`","data","config","<","data-config-label",">","`",".","optim","(",":","obj",":","`","Union","[","DictConfig",",","Dict","[","str",",","Any","]","]","`",",","`","optional","`",")",":","optimization","configuration",".","See","possible","options","in",":","ref",":","`","optimization","<","optimization-label",">","`","and","in","`","primer","<","https",":","\/\/github.com\/NVIDIA\/NeMo\/blob\/main\/tutorials\/00_NeMo_Primer.ipynb",">","`","_","tutorial","."]
code: def update_config_after_restoring_from_checkpoint(
      self, **kwargs
  ) -> None:
      """
      Set new values for some sections of config. Useful after restoring from checkpoint for fine-tuning
      and testing if config parameters of a restored checkpoint are not suitable.
  
      For ``class_labels``, ``common_dataset_parameters``, ``train_ds``, ``validation_ds``, ``test_ds``, there is
      no need to provide values for all items in an updated config section. If an item is omitted in this method
      parameter, then corresponding item in model config does not change.
  
      If the entire updated section is missing in the model config, then omitted items from this method parameters
      are set according to default values listed
      :ref:`here <run-config-label>`.
  
      .. warning::
          Parameter ``optim`` is processed in a special way. ``optim`` contents are used not for updating of
          model config, but for replacement of entire config section.
  
      If one of parameters ``train_ds``, ``validation_ds``, ``test_ds``, is provided but its value is
      ``None``, then corresponding section is replaced with ``None``.
  
      .. warning::
          You may change values of parameters related to label ids:
  
              - ``common_dataset_parameters.punct_label_ids``,
              - ``common_dataset_parameters.capit_label_ids``,
              - ``common_dataset_parameters.label_vocab_dir``,
              - ``class_labels.punct_labels_file``,
              - ``class_labels.capit_labels_file``,
  
          yet label ids in these parameters must be equal to label ids loaded from checkpoint. Otherwise,
          an error will be raised.
  
      Keyword Args:
          class_labels (:obj:`Union[DictConfig, Dict[str, str]]`): names of label id files used as label
              id dictionaries. See more in :ref:`class labels' config<class-labels-config-label>`.
          common_dataset_parameters (:obj:`Union[DictConfig, Dict[str, Any]]`, `optional`): see more in
              :ref:`common dataset parameters config<common-dataset-parameters-config-label>`.
          train_ds (:obj:`Union[DictConfig, Dict[str, Any]]`, `optional`): configuration of training dataset. See
              possible options in :ref:`data config<data-config-label>`.
          validation_ds (:obj:`Union[DictConfig, Dict[str, Any]]`, `optional`): configuration of validation
              dataset. See possible options in :ref:`data config<data-config-label>`.
          test_ds (:obj:`Union[DictConfig, Dict[str, Any]]`, `optional`): configuration of test dataset. See
              possible options in :ref:`data config<data-config-label>`.
          optim (:obj:`Union[DictConfig, Dict[str, Any]]`, `optional`): optimization configuration. See possible
              options in :ref:`optimization<optimization-label>` and in `primer
              <https://github.com/NVIDIA/NeMo/blob/main/tutorials/00_NeMo_Primer.ipynb>`_ tutorial.
      """
      allowed_keys = {
          "class_labels",
          "common_dataset_parameters",
          "train_ds",
          "validation_ds",
          "test_ds",
          "optim",
      }
      unexpected_keys = set(kwargs) - allowed_keys
      if unexpected_keys:
          raise ValueError(
              f"Found unexpected keyword arguments: {unexpected_keys}. You can use only {allowed_keys}."
          )
      if "class_labels" in kwargs:
          if kwargs["class_labels"] is None:
              raise ValueError(
                  f"'class_labels' parameters is `None`, whereas you cannot remove section 'class_labels' from model "
                  f"config."
              )
          self._cfg.class_labels = OmegaConf.merge(
              self._cfg.class_labels,
              OmegaConf.create(kwargs["class_labels"]),
          )
      if "common_dataset_parameters" in kwargs:
          if kwargs["common_dataset_parameters"] is None:
              raise ValueError(
                  f"'common_dataset_parameters' item is `None`, whereas you cannot remove section"
                  f"'common_dataset_parameters' from model config."
              )
          self._cfg.common_dataset_parameters = OmegaConf.merge(
              self._cfg.common_dataset_parameters,
              OmegaConf.create(kwargs["common_dataset_parameters"]),
          )
          self._check_label_config_parameters()
      if "train_ds" in kwargs:
          if kwargs["train_ds"] is None:
              self._cfg.train_ds = None
          else:
              if (
                  "train_ds" in self._cfg
                  and self._cfg.train_ds is not None
              ):
                  base = self._cfg.train_ds
              else:
                  base = OmegaConf.structured(
                      PunctuationCapitalizationTrainDataConfig
                  )
              self._cfg.train_ds = OmegaConf.merge(
                  base, OmegaConf.create(kwargs["train_ds"])
              )
      if "validation_ds" in kwargs:
          if kwargs["validation_ds"] is None:
              self._cfg.validation_ds = None
          else:
              if (
                  "validation_ds" in self._cfg
                  and self._cfg.validation_ds is not None
              ):
                  base = self._cfg.validation_ds
              else:
                  base = OmegaConf.structured(
                      PunctuationCapitalizationEvalDataConfig
                  )
              self._cfg.validation_ds = OmegaConf.merge(
                  base, OmegaConf.create(kwargs["validation_ds"])
              )
      if "test_ds" in kwargs:
          if kwargs["test_ds"] is None:
              self._cfg.test_ds = None
          else:
              if (
                  "test_ds" in self._cfg
                  and self._cfg.test_ds is not None
              ):
                  base = self._cfg.test_ds
              else:
                  base = OmegaConf.structured(
                      PunctuationCapitalizationEvalDataConfig
                  )
              self._cfg.test_ds = OmegaConf.merge(
                  base, OmegaConf.create(kwargs["test_ds"])
              )
      if "optim" in kwargs:
          self._cfg.optim = kwargs["optim"]
  
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 13 for calling the PunctuationCapitalizationModel obj's underlying member method code with example usage: obj.update_config_after_restoring_from_checkpoint() without return types