class_docstr: A dataset to use during training for punctuation and capitalization tasks.
  For inference, you will need
  :class:`~nemo.collections.nlp.data.token_classification.punctuation_capitalization_infer_dataset.BertPunctuationCapitalizationInferDataset`.
  For huge datasets which cannot be loaded into memory simultaneously use
  :class:`~nemo.collections.nlp.data.token_classification.punctuation_capitalization_tarred_dataset.BertPunctuationCapitalizationTarredDataset`.
  
  Args:
      text_file (:obj:`Union[str, os.PathLike]`): a path to a file with sequences, each line should contain a text
          without punctuation and capitalization
      labels_file (:obj:`Union[str, os.PathLike]`): a path to a file with labels, each line corresponds to word
          labels for a sentence in the ``text_file``. Labels have to follow format described in this section of
          documentation :ref:`NeMo Data Format<nemo-data-format-label>`.
      max_seq_length (:obj:`int`): max number of tokens in a source sequence. ``max_seq_length`` includes for [CLS]
          and [SEP] tokens. Sequences which are too long will be clipped by removal of tokens from the end of the
          sequence.
      tokenizer (:obj:`TokenizerSpec`): a tokenizer instance which has properties ``unk_id``, ``sep_id``, ``bos_id``,
          ``eos_id``.
      num_samples (:obj:`int`, `optional`, defaults to :obj:`-1`): a number of samples you want to use for the
          dataset. If ``-1``, use all dataset. Useful for testing.
      tokens_in_batch (:obj:`int`, `optional`, defaults to :obj:`5000`): number of tokens in a batch including
          paddings and special tokens ([CLS], [SEP], [UNK]). This class :meth:`__getitem__` method returns not
          samples but ready batches. Number of samples in a batch is adjusted for input sequences lengths. If input
          sequences are short, then a batch will contain more samples. Before packing into batches, samples are
          sorted by number of tokens they contain. Sorting allows to reduce number of pad tokens in a batch
          significantly. Regular PyTorch data loader shuffling will only permute batches with changing their content.
          Proper shuffling is achieved via calling method :meth:`repack_batches_with_shuffle` every epoch. If
          parameter ``number_of_batches_is_multiple_of`` is greater than 1, some batches may be split into smaller
          pieces.
      pad_label (:obj:`str`, `optional`, defaults to :obj:`'O'`): pad value to use for labels. It's also the neutral
          label both for punctuation and capitalization.
      punct_label_ids (:obj:`Dict[str, int]`, `optional`): dict to map punctuation labels to label ids. For dev set,
          use label ids generated during training to support cases when not all labels are present in the dev set.
          For training, it is recommended to set ``punct_label_ids`` to ``None`` or load from cache.
      capit_label_ids (:obj:`Dict[str, int]`, `optional`): same ``punct_label_ids`` for capitalization labels.
      ignore_extra_tokens (:obj:`bool`, `optional`, defaults to :obj:`False`): whether to compute loss on
          tokens which are not first tokens in a word. For example, assume that word ``'tokenization'`` is tokenized
          into ``['token', 'ization']``. If ``ignore_extra_tokens=True``, loss mask for the word is
          ``[True, False]``, and if ``ignore_extra_tokens=False``, then loss mask is ``[True, True]``.
      ignore_start_end (:obj:`bool`, `optional`, defaults to :obj:`True`): whether to ignore [CLS] and [SEP] tokens
          in the loss_mask.
      use_cache (:obj:`bool`, `optional`, defaults to :obj:`True`): whether to use pickled features already present
          in ``cache_dir`` or not. If pickled features file does not exist or ``use_cache=False``, then features are
          pickled in ``cache_dir``. Pickled features include input ids, subtokens mask (mask of first tokens in
          words), encoded punctuation and capitalization labels, label ids. Features creation consumes considerable
          time and this ``use_cache=True`` significantly speeds up training starting. Pickled features are also
          used for sharing features between processes if data parallel training is used.
      cache_dir (:obj:`Union[str, os.PathLike]`, `optional`): a path to a directory where cache (pickled features)
          is stored. By default, ``text_file`` parent directory is used. This parameter is useful if dataset
          directory is read-only, and you wish to pickle features. In such a case specify a path to directory which
          allows writing in ``cache_dir`` parameter.
      get_label_frequencies (:obj:`bool`, `optional`, defaults to :obj:`False`): whether to print and save label
          frequencies. Frequencies are showed if ``verbose`` parameter is ``True``. If
          ``get_label_frequencies=True``, then frequencies are saved into ``label_info_save_dir`` directory.
      label_info_save_dir (:obj:`Union[str, os.PathLike]`, `optional`): a path to a directory where label frequencies
          are saved. By default, a ``text_file`` parent directory is used. When method
          :meth:`save_labels_and_get_file_paths` is called label ids are saved into ``label_info_save_dir``
          directory. This parameter is useful if directory containing ``text_file`` is read-only.
      punct_label_vocab_file (:obj:`Union[str, os.PathLike]`, `optional`): a path to a .csv file containing
          punctuation label vocabulary. Each line in such a vocabulary file contains exactly one label. The first
          line has to contain `pad_label`, otherwise error will be raised.
      capit_label_vocab_file (:obj:`Union[str, os.PathLike]`, `optional`): same as ``punct_label_vocab_file`` for
          capitalization labels.
      add_masks_and_segment_ids_to_batch (:obj:`bool`, `optional`, defaults to :obj:`True`): whether to add
          ``'loss_mask'``, ``'input_mask'``, ``'segment_ids'`` items to a batch. Useful for creation of tarred
          dataset and can NOT be used during model training and inference.
      verbose (:obj:`bool`, `optional`, defaults to :obj:`True`): whether to show data examples, label stats and
          other useful information.
      n_jobs (:obj:`int`, `optional`, defaults to :obj:`0`): number of workers used for tokenization, encoding
          labels, creating "first token in word" mask, and clipping. If ``n_jobs <= 0`` data preparation is performed
          without multiprocessing. By default, ``n_jobs`` is ``0``.
  
          .. warning::
              There can be deadlocking problems with some tokenizers (e.g. SentencePiece, HuggingFace AlBERT)
              if ``n_jobs > 0``.
  
      number_of_batches_is_multiple_of (:obj:`int`, `optional`, defaults to :obj:`1`): number of batches in the
          dataset is made divisible by ``number_of_batches_is_multiple_of``. If ``number_of_batches_is_multiple_of``
          is greater than 1, then several batches are split in parts until number of batches
          is divisible by ``number_of_batches_is_multiple_of``. If there is no enough queries in the dataset to
          create enough batches, then a warning is printed. This parameter is useful for dev and validation datasets
          if multiple GPUs are used. The problem is that if number of batches is not evenly divisible by number of
          GPUs, then some queries may be processed several times and metrics will be distorted.
      batch_shuffling_random_seed (:obj:`int`, defaults to :obj:`int`): a random seed used for batches repacking and
          shuffling.
      tokenization_progress_queue (:obj:`multiprocessing.Queue`, `optional`): a queue for reporting tokenization
          progress. Useful for creation of tarred dataset
      batch_mark_up_progress_queue (:obj:`multiprocessing.Queue`, `optional`): a queue for reporting progress in
          deciding which samples batches will contain. Useful for creation of tarred dataset
      batch_building_progress_queue (:obj:`multiprocessing.Queue`, `optional`): a queue for reporting progress in
          batch creation (stacking and padding). Useful for creation of tarred dataset
      use_audio (:obj:`bool`, `optional`, defaults to :obj: `False`): If set to True dataset will return audio as well as text.
      audio_file (:obj:`Union[str, os.PathLike]`, `optional`): a path to file with audio paths.
      sample_rate (:obj:`int`, `optional`, defaults to :obj:`None`): sample rate of audios. Can be used for up sampling or down sampling of audio.
      use_bucketing (:obj:`bool`, `optional`, defaults to :obj: `True`): If set to False dataset will return ``batch_size`` batches instead of ``number_of_tokens`` tokens.
      preload_audios (:obj:`bool`, `optional`, defaults to :obj: `True`): If set to True batches will include waveforms, if set to False will store audio_filepaths instead and load audios during ``collate_fn`` call
class_docstr_tok: ["A","dataset","to","use","during","training","for","punctuation","and","capitalization","tasks",".","For","inference",",","you","will","need",":","class",":","`","~nemo.collections.nlp.data.token_classification.punctuation_capitalization_infer_dataset.BertPunctuationCapitalizationInferDataset","`",".","For","huge","datasets","which","can","not","be","loaded","into","memory","simultaneously","use",":","class",":","`","~nemo.collections.nlp.data.token_classification.punctuation_capitalization_tarred_dataset.BertPunctuationCapitalizationTarredDataset","`",".","Args",":","text_file","(",":","obj",":","`","Union","[","str",",","os.PathLike","]","`",")",":","a","path","to","a","file","with","sequences",",","each","line","should","contain","a","text","without","punctuation","and","capitalization","labels_file","(",":","obj",":","`","Union","[","str",",","os.PathLike","]","`",")",":","a","path","to","a","file","with","labels",",","each","line","corresponds","to","word","labels","for","a","sentence","in","the","``","text_file","``",".","Labels","have","to","follow","format","described","in","this","section","of","documentation",":","ref",":","`","NeMo","Data","Format","<","nemo-data-format-label",">","`",".","max_seq_length","(",":","obj",":","`","int","`",")",":","max","number","of","tokens","in","a","source","sequence",".","``","max_seq_length","``","includes","for","[","CLS","]","and","[","SEP","]","tokens",".","Sequences","which","are","too","long","will","be","clipped","by","removal","of","tokens","from","the","end","of","the","sequence",".","tokenizer","(",":","obj",":","`","TokenizerSpec","`",")",":","a","tokenizer","instance","which","has","properties","``","unk_id","``",",","``","sep_id","``",",","``","bos_id","``",",","``","eos_id","``",".","num_samples","(",":","obj",":","`","int","`",",","`","optional","`",",","defaults","to",":","obj",":","`","-1","`",")",":","a","number","of","samples","you","want","to","use","for","the","dataset",".","If","``","-1","``",",","use","all","dataset",".","Useful","for","testing",".","tokens_in_batch","(",":","obj",":","`","int","`",",","`","optional","`",",","defaults","to",":","obj",":","`","5000","`",")",":","number","of","tokens","in","a","batch","including","paddings","and","special","tokens","(","[","CLS","]",",","[","SEP","]",",","[","UNK","]",")",".","This","class",":","meth",":","`","__getitem__","`","method","returns","not","samples","but","ready","batches",".","Number","of","samples","in","a","batch","is","adjusted","for","input","sequences","lengths",".","If","input","sequences","are","short",",","then","a","batch","will","contain","more","samples",".","Before","packing","into","batches",",","samples","are","sorted","by","number","of","tokens","they","contain",".","Sorting","allows","to","reduce","number","of","pad","tokens","in","a","batch","significantly",".","Regular","PyTorch","data","loader","shuffling","will","only","permute","batches","with","changing","their","content",".","Proper","shuffling","is","achieved","via","calling","method",":","meth",":","`","repack_batches_with_shuffle","`","every","epoch",".","If","parameter","``","number_of_batches_is_multiple_of","``","is","greater","than","1",",","some","batches","may","be","split","into","smaller","pieces",".","pad_label","(",":","obj",":","`","str","`",",","`","optional","`",",","defaults","to",":","obj",":","`","'","O","'","`",")",":","pad","value","to","use","for","labels",".","It","'s","also","the","neutral","label","both","for","punctuation","and","capitalization",".","punct_label_ids","(",":","obj",":","`","Dict","[","str",",","int","]","`",",","`","optional","`",")",":","dict","to","map","punctuation","labels","to","label","ids",".","For","dev","set",",","use","label","ids","generated","during","training","to","support","cases","when","not","all","labels","are","present","in","the","dev","set",".","For","training",",","it","is","recommended","to","set","``","punct_label_ids","``","to","``","None","``","or","load","from","cache",".","capit_label_ids","(",":","obj",":","`","Dict","[","str",",","int","]","`",",","`","optional","`",")",":","same","``","punct_label_ids","``","for","capitalization","labels",".","ignore_extra_tokens","(",":","obj",":","`","bool","`",",","`","optional","`",",","defaults","to",":","obj",":","`","False","`",")",":","whether","to","compute","loss","on","tokens","which","are","not","first","tokens","in","a","word",".","For","example",",","assume","that","word","``","'tokenization","'","``","is","tokenized","into","``","[","'token","'",",","'ization","'","]","``",".","If","``","ignore_extra_tokens=True","``",",","loss","mask","for","the","word","is","``","[","True",",","False","]","``",",","and","if","``","ignore_extra_tokens=False","``",",","then","loss","mask","is","``","[","True",",","True","]","``",".","ignore_start_end","(",":","obj",":","`","bool","`",",","`","optional","`",",","defaults","to",":","obj",":","`","True","`",")",":","whether","to","ignore","[","CLS","]","and","[","SEP","]","tokens","in","the","loss_mask",".","use_cache","(",":","obj",":","`","bool","`",",","`","optional","`",",","defaults","to",":","obj",":","`","True","`",")",":","whether","to","use","pickled","features","already","present","in","``","cache_dir","``","or","not",".","If","pickled","features","file","does","not","exist","or","``","use_cache=False","``",",","then","features","are","pickled","in","``","cache_dir","``",".","Pickled","features","include","input","ids",",","subtokens","mask","(","mask","of","first","tokens","in","words",")",",","encoded","punctuation","and","capitalization","labels",",","label","ids",".","Features","creation","consumes","considerable","time","and","this","``","use_cache=True","``","significantly","speeds","up","training","starting",".","Pickled","features","are","also","used","for","sharing","features","between","processes","if","data","parallel","training","is","used",".","cache_dir","(",":","obj",":","`","Union","[","str",",","os.PathLike","]","`",",","`","optional","`",")",":","a","path","to","a","directory","where","cache","(","pickled","features",")","is","stored",".","By","default",",","``","text_file","``","parent","directory","is","used",".","This","parameter","is","useful","if","dataset","directory","is","read-only",",","and","you","wish","to","pickle","features",".","In","such","a","case","specify","a","path","to","directory","which","allows","writing","in","``","cache_dir","``","parameter",".","get_label_frequencies","(",":","obj",":","`","bool","`",",","`","optional","`",",","defaults","to",":","obj",":","`","False","`",")",":","whether","to","print","and","save","label","frequencies",".","Frequencies","are","showed","if","``","verbose","``","parameter","is","``","True","``",".","If","``","get_label_frequencies=True","``",",","then","frequencies","are","saved","into","``","label_info_save_dir","``","directory",".","label_info_save_dir","(",":","obj",":","`","Union","[","str",",","os.PathLike","]","`",",","`","optional","`",")",":","a","path","to","a","directory","where","label","frequencies","are","saved",".","By","default",",","a","``","text_file","``","parent","directory","is","used",".","When","method",":","meth",":","`","save_labels_and_get_file_paths","`","is","called","label","ids","are","saved","into","``","label_info_save_dir","``","directory",".","This","parameter","is","useful","if","directory","containing","``","text_file","``","is","read-only",".","punct_label_vocab_file","(",":","obj",":","`","Union","[","str",",","os.PathLike","]","`",",","`","optional","`",")",":","a","path","to","a",".csv","file","containing","punctuation","label","vocabulary",".","Each","line","in","such","a","vocabulary","file","contains","exactly","one","label",".","The","first","line","has","to","contain","`","pad_label","`",",","otherwise","error","will","be","raised",".","capit_label_vocab_file","(",":","obj",":","`","Union","[","str",",","os.PathLike","]","`",",","`","optional","`",")",":","same","as","``","punct_label_vocab_file","``","for","capitalization","labels",".","add_masks_and_segment_ids_to_batch","(",":","obj",":","`","bool","`",",","`","optional","`",",","defaults","to",":","obj",":","`","True","`",")",":","whether","to","add","``","'loss_mask","'","``",",","``","'input_mask","'","``",",","``","'segment_ids","'","``","items","to","a","batch",".","Useful","for","creation","of","tarred","dataset","and","can","NOT","be","used","during","model","training","and","inference",".","verbose","(",":","obj",":","`","bool","`",",","`","optional","`",",","defaults","to",":","obj",":","`","True","`",")",":","whether","to","show","data","examples",",","label","stats","and","other","useful","information",".","n_jobs","(",":","obj",":","`","int","`",",","`","optional","`",",","defaults","to",":","obj",":","`","0","`",")",":","number","of","workers","used","for","tokenization",",","encoding","labels",",","creating","``","first","token","in","word","''","mask",",","and","clipping",".","If","``","n_jobs","<","=","0","``","data","preparation","is","performed","without","multiprocessing",".","By","default",",","``","n_jobs","``","is","``","0","``",".","..","warning",":",":","There","can","be","deadlocking","problems","with","some","tokenizers","(","e.g",".","SentencePiece",",","HuggingFace","AlBERT",")","if","``","n_jobs",">","0","``",".","number_of_batches_is_multiple_of","(",":","obj",":","`","int","`",",","`","optional","`",",","defaults","to",":","obj",":","`","1","`",")",":","number","of","batches","in","the","dataset","is","made","divisible","by","``","number_of_batches_is_multiple_of","``",".","If","``","number_of_batches_is_multiple_of","``","is","greater","than","1",",","then","several","batches","are","split","in","parts","until","number","of","batches","is","divisible","by","``","number_of_batches_is_multiple_of","``",".","If","there","is","no","enough","queries","in","the","dataset","to","create","enough","batches",",","then","a","warning","is","printed",".","This","parameter","is","useful","for","dev","and","validation","datasets","if","multiple","GPUs","are","used",".","The","problem","is","that","if","number","of","batches","is","not","evenly","divisible","by","number","of","GPUs",",","then","some","queries","may","be","processed","several","times","and","metrics","will","be","distorted",".","batch_shuffling_random_seed","(",":","obj",":","`","int","`",",","defaults","to",":","obj",":","`","int","`",")",":","a","random","seed","used","for","batches","repacking","and","shuffling",".","tokenization_progress_queue","(",":","obj",":","`","multiprocessing.Queue","`",",","`","optional","`",")",":","a","queue","for","reporting","tokenization","progress",".","Useful","for","creation","of","tarred","dataset","batch_mark_up_progress_queue","(",":","obj",":","`","multiprocessing.Queue","`",",","`","optional","`",")",":","a","queue","for","reporting","progress","in","deciding","which","samples","batches","will","contain",".","Useful","for","creation","of","tarred","dataset","batch_building_progress_queue","(",":","obj",":","`","multiprocessing.Queue","`",",","`","optional","`",")",":","a","queue","for","reporting","progress","in","batch","creation","(","stacking","and","padding",")",".","Useful","for","creation","of","tarred","dataset","use_audio","(",":","obj",":","`","bool","`",",","`","optional","`",",","defaults","to",":","obj",":","`","False","`",")",":","If","set","to","True","dataset","will","return","audio","as","well","as","text",".","audio_file","(",":","obj",":","`","Union","[","str",",","os.PathLike","]","`",",","`","optional","`",")",":","a","path","to","file","with","audio","paths",".","sample_rate","(",":","obj",":","`","int","`",",","`","optional","`",",","defaults","to",":","obj",":","`","None","`",")",":","sample","rate","of","audios",".","Can","be","used","for","up","sampling","or","down","sampling","of","audio",".","use_bucketing","(",":","obj",":","`","bool","`",",","`","optional","`",",","defaults","to",":","obj",":","`","True","`",")",":","If","set","to","False","dataset","will","return","``","batch_size","``","batches","instead","of","``","number_of_tokens","``","tokens",".","preload_audios","(",":","obj",":","`","bool","`",",","`","optional","`",",","defaults","to",":","obj",":","`","True","`",")",":","If","set","to","True","batches","will","include","waveforms",",","if","set","to","False","will","store","audio_filepaths","instead","and","load","audios","during","``","collate_fn","``","call"]
code: def _check_label_ids_loaded_from_pkl(
      self,
      parameter_punct_label_ids: Dict[str, int],
      parameter_capit_label_ids: Dict[str, int],
      pkl_punct_label_ids: Any,
      pkl_capit_label_ids: Any,
      punct_label_vocab_file: Optional[Path],
      capit_label_vocab_file: Optional[Path],
  ) -> None:
      if not isinstance(pkl_punct_label_ids, dict):
          raise ValueError(
              f"Punctuation label ids loaded from features file {self.features_pkl} have wrong type "
              f"{type(pkl_punct_label_ids)}"
          )
      if parameter_punct_label_ids is not None:
          if parameter_punct_label_ids != pkl_punct_label_ids:
              raise_not_equal_labels_error(
                  first_labels=parameter_punct_label_ids,
                  second_labels=pkl_punct_label_ids,
                  first_labels_desc="Punctuation labels passed in parameter `punct_label_ids`"
                  if punct_label_vocab_file is None
                  else f"Punctuation labels loaded from file {punct_label_vocab_file}",
                  second_labels_desc=f"Punctuation label ids loaded from features file {self.features_pkl}",
              )
      if not isinstance(pkl_capit_label_ids, dict):
          raise ValueError(
              f"Capitalization label ids loaded from features file {self.features_pkl} has wrong type "
              f"{type(pkl_capit_label_ids)}"
          )
      if parameter_capit_label_ids is not None:
          if parameter_capit_label_ids != pkl_capit_label_ids:
              raise_not_equal_labels_error(
                  first_labels=parameter_capit_label_ids,
                  second_labels=pkl_capit_label_ids,
                  first_labels_desc="Capitalization labels passed in parameter `capit_label_ids`"
                  if capit_label_vocab_file is None
                  else f"Capitalization labels loaded from file {capit_label_vocab_file}",
                  second_labels_desc=f"Capitalization label ids loaded from features file {self.features_pkl}",
              )
  
imports: import itertools
  import multiprocessing
  import os
  import pickle
  import tempfile
  from dataclasses import dataclass
  from math import ceil
  from pathlib import Path
  from queue import Empty
  from time import sleep
  from typing import Any, Dict, List, Optional, Set, Tuple, Union
  import numpy
  import torch
  from numpy import ndarray
  from omegaconf import MISSING, DictConfig, OmegaConf
  from torch.nn.utils.rnn import pad_sequence
  from tqdm import tqdm
  from nemo.collections.common.tokenizers.tokenizer_spec import TokenizerSpec
  from nemo.collections.nlp.data.data_utils.data_preprocessing import get_label_stats, get_stats
  from nemo.core.classes import Dataset
  from nemo.core.neural_types import AudioSignal, ChannelType, LabelsType, LengthsType, MaskType, NeuralType
  from nemo.utils import logging
  from nemo.utils.get_rank import is_global_rank_zero
label_desc: Use image node_id 5 for calling the BertPunctuationCapitalizationDataset obj's underlying member method code with example usage: obj._check_label_ids_loaded_from_pkl(parameter_punct_label_ids, parameter_capit_label_ids, pkl_punct_label_ids, pkl_capit_label_ids, punct_label_vocab_file, capit_label_vocab_file) without return types