o }oi @sddlZddlZddlZddlZddlZddlmZddlm Z m Z mZm Z m Z m Z mZddlZddlZddlZddlmZddlmZddlmZddlmZddlmZdd lmZdd lmZm Z dd l!m"Z"m#Z#dd l$Tdd l%m&Z&ddl%m'Z(ddl)m*Z*m+Z+m,Z,m-Z-ddl.m/Z/ddl0m1Z1ddl2m3Z3gdZ4d5gddde6DZ7ddZ8GdddZ9de:de;de;fdd Z< !   "dAd#ee:e e:fd$e=d%e e=d&e e;d'e;f d(d)Z> d#ee:e e:fde:d*e=de;de;f d+d,Z?Gd-d.d.e"Z@Gd/d0d0e@ZAGd1d2d2e@ZBe/d3d4Gd5d6d6e#ZCGd7d8d8eCZDGd9d:d:eCZEGd;d<dd>ZGGd?d@d@eZHdS)BN)Iterable)CallableDictrListOptionalTupleUnion) ChainDataset)tqdm)WaveformFeaturizer)ChannelSelectorType)available_formats) tokenizers) collectionsparsers)DatasetIterableDataset)*)logging) webdataset)DataStoreObjectdatastore_object_getis_datastore_cache_sharedis_datastore_path) deprecated)webdataset_split_by_workers)is_global_rank_zero)AudioToCharDatasetAudioToBPEDatasetTarredAudioToCharDatasetTarredAudioToBPEDataset;)wavmp3flacopuscCsg|]}|qS)lower).0fmtr&r&[/home/ubuntu/.local/lib/python3.10/site-packages/nemo/collections/asr/data/audio_to_text.py 1r+cCstt|}t|dkr|\}}}}}nt|dkr#d}|\}}}}ntdd}|ddu}|r7t|}|ddu} | rEt|} gg} } |D]R} t| dkr\| \}}}}}n| \}}}}|r|}||krzd||f}tjj ||}| || r|}|| krd| |f}tjjj |||d}| |qL|rt | } t |}nd\} }| rt | } t |}nd} d}|dur| || |fStj |tj d}| || ||fS) arcollate batch of audio sig, audio len, tokens, tokens len Args: batch (Optional[FloatTensor], Optional[LongTensor], LongTensor, LongTensor): A tuple of tuples of signal, signal lengths, encoded tokens, and encoded tokens length. This collate func assumes the signals are 1d torch tensors (i.e. mono audio). Nz$Expects 4 or 5 tensors in the batch!r)value)NNdtype)listziplen ValueErrormaxitemtorchnn functionalpadappendstacktensorint32)batchpad_id packed_batch_ audio_lengthstokens_lengths sample_ids max_audio_len has_audio has_tokensmax_tokens_len audio_signaltokensbsigsig_lentokens_i tokens_i_lenr;r&r&r*_speech_collate_fn4sZ                   rRc@seZdZdZ        ddedeeefdeedeed e d ee d ee d e d e deefddZ de de e e e ffddZdede e e e ffddZdejjde e e e ffddZdS)ASRManifestProcessora Class that processes a manifest json file containing paths to audio files, transcripts, and durations (in seconds). Each new line is a different sample. Example below: {"audio_filepath": "/path/to/audio.wav", "text_filepath": "/path/to/audio.txt", "duration": 23.147} ... {"audio_filepath": "/path/to/audio.wav", "text": "the transcription", "offset": 301.75, "duration": 0.82, "utt": "utterance_id", "ctm_utt": "en_4156", "side": "A"} Args: manifest_filepath: Path to manifest json as described above. Can be comma-separated paths. parser: Str for a language specific preprocessor or a callable. max_duration: If audio exceeds this length, do not include in dataset. min_duration: If audio is less than this length, do not include in dataset. max_utts: Limit number of utterances. bos_id: Id of beginning of sequence symbol to append if not None. eos_id: Id of end of sequence symbol to append if not None. pad_id: Id of pad symbol. Defaults to 0. NrFmanifest_filepathparser max_duration min_durationmax_uttsbos_ideos_idrAindex_by_file_idmanifest_parse_funcc Cs6||_tj|||||| | d|_||_||_||_dS)N)manifests_filesrUrWrV max_numberr[ parse_func)rUr ASRAudioText collectionrZrYrA) selfrTrUrVrWrXrYrZrAr[r\r&r&r*__init__s   zASRManifestProcessor.__init__indexreturncCs|j|}||SN)raprocess_text_by_sample)rbrdsampler&r&r*process_text_by_ids  z'ASRManifestProcessor.process_text_by_idfile_idcCs$|jj|d}|j|}||SNr)ramappingrg)rbrj manifest_idxrhr&r&r*process_text_by_file_ids  z,ASRManifestProcessor.process_text_by_file_idrhcCsV|jt|j}}|jdur|jg|}|d7}|jdur'||jg}|d7}||fS)N) text_tokensr4rYrZ)rbrhttlr&r&r*rgs    z+ASRManifestProcessor.process_text_by_sample)NNrNNrFN)__name__ __module__ __qualname____doc__strrrrfloatintboolrcrrrirnrr` OUTPUT_TYPErgr&r&r&r*rSqsD    &rSshard_strategy world_size global_rankc Cs>ddg}||vrtd|t|tr:gd}|D] }||vr&||d}qgd}|D] }||vr9||d}q-t|trHttj|dd }|d kr|dkrtd t||d krkt d t|d|dt|||}|t||} ||| }td||| |S|dkrtd|Std||S)Nscatter replicatez `shard_strategy` must be one of )([<_OP_{))]>_CL_}F)escaperozDAll tarred dataset shards will be scattered evenly across all nodes.rz$Number of shards in tarred dataset (z5) is not divisible by number of distributed workers (z).z@Partitioning tarred dataset: process (%d) taking shards [%d, %d)z>All tarred dataset shards will be replicated across all nodes.z.Invalid shard strategy ! Allowed values are : ) r5 isinstancerwreplacer2 braceexpandrinfor4warning) sharded_filepathsr|r}r~valid_shard_strategiesbrace_keys_openbkeybrace_keys_close begin_idxend_idxr&r&r*expand_sharded_filepathssJ         rFmanifest_filepaths cache_audio shared_cache num_workersmax_num_workersc Cs t|tr |d}tdd|D}|dkrdd}tjrntjrnt d|dur1t }|r.rc Ss|dur td}t||}|D]}t|rtd|t|}tdt||rtd|tj |}g}t |d#}|D]} t | } tj || d} |t| dqCWdn1sfwY|dur|dkrtd |tj|d } tt| t|t|d } Wdn1swYntd g} t|D] }| |duqt| std tdqtd|qdS)z1Cache manifests and audio data from object store.NrozCache manifest file: %sz Cached at: %sz"Cache audio from manifest file: %sraudio_filepath) store_pathz+Using multiprocessing with num_workers: %d.) processes)totalzUsing a single process.z&Some files not downloaded successfullyzCaching completez#Manifest is not on a data store: %s)os cpu_countminrrrrgetrwpathdirnameopenjsonloadsjoinr<debugmultiprocessingPoolr2r imaprr4all RuntimeError)rrrr manifest_filecached_manifest_file manifest_dir audio_objectsrliner7rpresult audio_objectr&r&r* cache_datasJ            z-cache_datastore_manifests..cache_dataz5Distributed environment is available and initialized.Nz.)r IterableABCr)rbrdr&rr* __getitem__s  z_AudioTextDataset.__getitem__c Cs|jj|}|j}|durd}|jj|j||j|j|j|j d}|t |j d }}|jj|d\}}|jrL||t | t | |f} | S||t | t | f} | S)Nr)offsetdurationrorig_srrrh)rrarrprocess audio_filerrrrr8r>shapelongrgr) rbrdrhrfeaturesrflrqrroutputr&r&r*rs& " z!_AudioTextDataset._process_samplecCs t|jjSrf)r4rrarr&r&r*__len__rz_AudioTextDataset.__len__cCst||jjdS)N)rA)rRrrArbr@r&r&r* _collate_fnsz_AudioTextDataset._collate_fn) FNNNrFNNrFNN)rsrtrurvpropertyrrrwrrrrryrzr rcrrrrrr&r&r&r*rsh       ( rc(seZdZdZedeeeeffddZ              d!d ed e ee efd e de dddeedeede de de de de dee dee de de eefde deedeef&fdd ZZS)"ra Dataset that loads tensors via a json file containing paths to audio files, transcripts, and durations (in seconds). Each new line is a different sample. Example below: {"audio_filepath": "/path/to/audio.wav", "text_filepath": "/path/to/audio.txt", "duration": 23.147} ... {"audio_filepath": "/path/to/audio.wav", "text": "the transcription", "offset": 301.75, "duration": 0.82, "utt": "utterance_id", "ctm_utt": "en_4156", "side": "A"} Args: manifest_filepath: Path to manifest json as described above. Can be comma-separated paths. labels: String containing all the possible characters to map to sample_rate (int): Sample rate to resample loaded audio to int_values (bool): If true, load samples as 32-bit integers. Defauts to False. augmentor (nemo.collections.asr.parts.perturb.AudioAugmentor): An AudioAugmentor object used to augment loaded audio max_duration: If audio exceeds this length, do not include in dataset min_duration: If audio is less than this length, do not include in dataset max_utts: Limit number of utterances blank_index: blank character index, default = -1 unk_index: unk_character index, default = -1 normalize: whether to normalize transcript text (default): True bos_id: Id of beginning of sequence symbol to append if not None eos_id: Id of end of sequence symbol to append if not None return_sample_id (bool): whether to return the sample_id as a part of each sample channel_selector (int | Iterable[int] | str): select a single channel or a subset of channels from multi-channel audio. If set to `'average'`, it performs averaging across channels. Disabled if set to `None`. Defaults to `None`. Uses zero-based indexing. manifest_parse_func: Optional function to parse manifest entries. Defaults to None. rec Crrrrr&r&r*r"rzAudioToCharDataset.output_typesFNrTenrTlabelsrrrrrVrWrX blank_index unk_index normalizerrYrZrArUrrr\csH||_tj||| | | d}tj||||||||| | |||||ddS)Nrnameunk_idblank_id do_normalize)rTrUrrrrVrWrXrrYrZrArrr\rr make_parsersuperrc)rbrTrrrrrVrWrXr r r rrYrZrArUrrr\ __class__r&r*rc-s*  zAudioToCharDataset.__init__)FNNNrrrTFNNrrFNN)rsrtrurvrrrrwrrrrryrzrxrr rc __classcell__r&r&rr*rsv!      rcseZdZdZedeeeeffddZ          dd ed d d e d e dddee dee de de de de dee dee ffdd ZZS)ra} Dataset that loads tensors via a json file containing paths to audio files, transcripts, and durations (in seconds). Each new line is a different sample. Example below: {"audio_filepath": "/path/to/audio.wav", "text_filepath": "/path/to/audio.txt", "duration": 23.147} ... {"audio_filepath": "/path/to/audio.wav", "text": "the transcription", "offset": 301.75, "duration": 0.82, "utt": "utterance_id", "ctm_utt": "en_4156", "side": "A"} In practice, the dataset and manifest used for character encoding and byte pair encoding are exactly the same. The only difference lies in how the dataset tokenizes the text in the manifest. Args: manifest_filepath: Path to manifest json as described above. Can be comma-separated paths. tokenizer: A subclass of the Tokenizer wrapper found in the common collection, nemo.collections.common.tokenizers.TokenizerSpec. ASR Models support a subset of all available tokenizers. sample_rate (int): Sample rate to resample loaded audio to int_values (bool): If true, load samples as 32-bit integers. Defauts to False. augmentor (nemo.collections.asr.parts.perturb.AudioAugmentor): An AudioAugmentor object used to augment loaded audio max_duration: If audio exceeds this length, do not include in dataset min_duration: If audio is less than this length, do not include in dataset max_utts: Limit number of utterances trim: Whether to trim silence segments use_start_end_token: Boolean which dictates whether to add [BOS] and [EOS] tokens to beginning and ending of speech respectively. return_sample_id (bool): whether to return the sample_id as a part of each sample channel_selector (int | Iterable[int] | str): select a single channel or a subset of channels from multi-channel audio. If set to `'average'`, it performs averaging across channels. Disabled if set to `None`. Defaults to `None`. Uses zero-based indexing. manifest_parse_func: Optional function to parse manifest entries. Defaults to None. rec Crrrrr&r&r*rrzAudioToBPEDataset.output_typesFNrTrT tokenizer0nemo.collections.common.tokenizers.TokenizerSpecrrrrrVrWrXruse_start_end_tokenrrr\cs| rt|dr|jdkr|j}nd}| r"t|dr"|jdkr"|j}nd}t|dr2|jdkr2|j}nd}Gddd}tj||||||||||||| | | | ddS)NrYrrZrAc@eZdZddZddZdS)z4AudioToBPEDataset.__init__..TokenizerWrappercS&t|tjjr d|_nd|_||_dSNTFrraggregate_tokenizerAggregateTokenizer is_aggregate _tokenizerrbrr&r&r*rc z=AudioToBPEDataset.__init__..TokenizerWrapper.__init__cWVt|dtr#|jr#g}|dD]}||j|d|dq|S|jj|}|SNrrwlangrrr extendr! text_to_idsrbargsrqspanr&r&r*__call__  z=AudioToBPEDataset.__init__..TokenizerWrapper.__call__Nrsrtrurcr-r&r&r&r*TokenizerWrapper r0)rTrUrrrrVrWrXrYrZrArrrr\hasattrrYrZrArrc)rbrTrrrrrVrWrXrrrrr\rYrZrAr0rr&r*rcs6 zAudioToBPEDataset.__init__) FNNNrFTFNN)rsrtrurvrrrrwrrryrzr rrcrr&r&rr*r\sR%     rzlWebdataset support will be removed in v2.1.0 versions, please use LhotseSpeechToTextBpeDataset class instead) explanationc(@seZdZdZ               d,deeeefdeded ed e d e d d ede e de e de de ede ededede dedede de ef&ddZ ddZ ddZd d!Zd"d#Zd$d%Zd&d'Zd(d)Zd*d+ZdS)-_TarredAudioToTextDataseta A similar Dataset to the AudioToCharDataset/AudioToBPEDataset, but which loads tarred audio files. Accepts a single comma-separated JSON manifest file (in the same style as for the AudioToCharDataset/AudioToBPEDataset), as well as the path(s) to the tarball(s) containing the wav files. Each line of the manifest should contain the information for one audio file, including at least the transcript and name of the audio file within the tarball. Valid formats for the audio_tar_filepaths argument include: (1) a single string that can be brace-expanded, e.g. 'path/to/audio.tar' or 'path/to/audio_{1..100}.tar.gz', or (2) a list of file paths that will not be brace-expanded, e.g. ['audio_1.tar', 'audio_2.tar', ...]. Note: For brace expansion in (1), there may be cases where `{x..y}` syntax cannot be used due to shell interference. This occurs most commonly inside SLURM scripts. Therefore we provide a few equivalent replacements. Supported opening braces - { <=> (, [, < and the special tag _OP_. Supported closing braces - } <=> ), ], > and the special tag _CL_. For SLURM based tasks, we suggest the use of the special tags for ease of use. See the WebDataset documentation for more information about accepted data and input formats. If using multiple workers the number of shards should be divisible by world_size to ensure an even split among workers. If it is not divisible, logging will give a warning but training will proceed. In addition, if using mutiprocessing, each shard MUST HAVE THE SAME NUMBER OF ENTRIES after filtering is applied. We currently do not check for this, but your program may hang if the shards are uneven! Notice that a few arguments are different from the AudioToCharDataset; for example, shuffle (bool) has been replaced by shuffle_n (int). Additionally, please note that the len() of this DataLayer is assumed to be the length of the manifest after filtering. An incorrect manifest length may lead to some DataLoader issues down the line. Args: audio_tar_filepaths: Either a list of audio tarball filepaths, or a string (can be brace-expandable). manifest_filepath (str): Path to the manifest. parser (callable): A callable which is used to pre-process the text output. sample_rate (int): Sample rate to resample loaded audio to int_values (bool): If true, load samples as 32-bit integers. Defauts to False. augmentor (nemo.collections.asr.parts.perturb.AudioAugmentor): An AudioAugmentor object used to augment loaded audio shuffle_n (int): How many samples to look ahead and load to be shuffled. See WebDataset documentation for more details. Defaults to 0. min_duration (float): Dataset parameter. All training files which have a duration less than min_duration are dropped. Note: Duration is read from the manifest JSON. Defaults to 0.1. max_duration (float): Dataset parameter. All training files which have a duration more than max_duration are dropped. Note: Duration is read from the manifest JSON. Defaults to None. blank_index (int): Blank character index, defaults to -1. unk_index (int): Unknown character index, defaults to -1. normalize (bool): Dataset parameter. Whether to use automatic text cleaning. It is highly recommended to manually clean text for best results. Defaults to True. trim (bool): Whether to use trim silence from beginning and end of audio signal using librosa.effects.trim(). Defaults to False. bos_id (id): Dataset parameter. Beginning of string symbol id used for seq2seq models. Defaults to None. eos_id (id): Dataset parameter. End of string symbol id used for seq2seq models. Defaults to None. pad_id (id): Token used to pad when collating samples in batches. If this is None, pads using 0s. Defaults to None. shard_strategy (str): Tarred dataset shard distribution strategy chosen as a str value during ddp. - `scatter`: The default shard strategy applied by WebDataset, where each node gets a unique set of shards, which are permanently pre-allocated and never changed at runtime. - `replicate`: Optional shard strategy, where each node gets all of the set of shards available in the tarred dataset, which are permanently pre-allocated and never changed at runtime. The benefit of replication is that it allows each node to sample data points from the entire dataset independently of other nodes, and reduces dependence on value of `shuffle_n`. .. warning:: Replicated strategy allows every node to sample the entire set of available tarfiles, and therefore more than one node may sample the same tarfile, and even sample the same data points! As such, there is no assured guarantee that all samples in the dataset will be sampled at least once during 1 epoch. Scattered strategy, on the other hand, on specific occasions (when the number of shards is not divisible with ``world_size``), will not sample the entire dataset. For these reasons it is not advisable to use tarred datasets as validation or test datasets. shard_manifests (bool): Whether or not to try / shard manifests. Defaults to False. global_rank (int): Worker rank, used for partitioning shards. Defaults to 0. world_size (int): Total number of processes, used for partitioning shards. Defaults to 0. return_sample_id (bool): whether to return the sample_id as a part of each sample manifest_parse_func: Optional function to parse manifest entries. Defaults to None. FNrraudio_tar_filepathsrTrUrrrr shuffle_nrWrVrrYrZrAr|rr~r}rr\c Cs||_t|||||d}t|dt||| |d| | | d|d |_||_t|||d|_| |_ | |_ | |_ | |_ ||_ t||||d}ttj|dtt|ttjtd d td d |j|jt|j |_dS) N)rr|rr}r~)rrT) rTrUrVrWrXrYrZrAr[r\rr)urls__key__)audiokeyr:r;)rrrrSr _compute_lenr4r rrrZrYrArrwds DataPipelineSimpleShardListrshuffletarfile_to_samplesrenameVALID_FILE_FORMATSto_tuple_filter _loop_offsetsmap _build_sample_dataset)rbr6rTrUrrrr7rWrVrrYrZrAr|rr~r}rr\r&r&r*rc1sZ       z"_TarredAudioToTextDataset.__init__cGfddd}||jjS)aThis function is used to remove samples that have been filtered out by ASRAudioText already. Otherwise, we would get a KeyError as _build_sample attempts to find the manifest entry for a sample that was filtered out (e.g. for duration). Note that if using multi-GPU training, filtering may lead to an imbalance in samples in each shard, which may make your code hang as one process will finish before the other. c(eZdZfddZddZddZdS)z<_TarredAudioToTextDataset._filter..TarredAudioFiltercs|_||_dSrf)iteratorrarbrarLr&r*rcs zE_TarredAudioToTextDataset._filter..TarredAudioFilter.__init__cS|Srfr&rr&r&r*__iter__zE_TarredAudioToTextDataset._filter..TarredAudioFilter.__iter__cSs> t|j\}}tjtj|\}}||jjvr||fSqrf)nextrLrrsplitextbasenamerarl)rb audio_bytesaudio_filenamerjrCr&r&r*__next__s  zE_TarredAudioToTextDataset._filter..TarredAudioFilter.__next__NrsrtrurcrPrWr&rNr&r*TarredAudioFilters  rYr)rbrLrYr&rNr*rEs z!_TarredAudioToTextDataset._filtercrJ)zYThis function is used to iterate through utterances with different offsets for each file.crK)zG_TarredAudioToTextDataset._loop_offsets..TarredAudioLoopOffsetscs"|_||_d|_d|_d|_dSrk)rLra current_fn current_bytes offset_idrMrNr&r*rcs  zP_TarredAudioToTextDataset._loop_offsets..TarredAudioLoopOffsets.__init__cSrOrfr&rr&r&r*rPrQzP_TarredAudioToTextDataset._loop_offsets..TarredAudioLoopOffsets.__iter__cSs||jdurt|j\|_|_d|_n$|jj|j}t||jdkr/t|j\|_|_d|_n|jd7_|j|j|jfS)Nrro)rZrRrLr[r\rarlr4)rb offset_listr&r&r*rWs zP_TarredAudioToTextDataset._loop_offsets..TarredAudioLoopOffsets.__next__NrXr&rNr&r*TarredAudioLoopOffsetss  r^r)rbrLr^r&rNr*rFs z'_TarredAudioToTextDataset._loop_offsetscCs t||jSrf)rRrArr&r&r*rrz%_TarredAudioToTextDataset._collate_fncCs>|\}}}tjtj|\}}|jjj||}|jj|}|j} | dur)d} t |} |j j | | |j |j |jd} | | t| jd} } |jt|j}}|jj|d|jdurl|jg|}|d7}|jdur{||jg}|d7}|jr| | t|t||fS| | t|t|fS)z\Builds the training sample by combining the data from the WebDataset with the manifest info.Nr)rrrrrro)rrrSrTrrarlrioBytesIOrrrrrcloser8r>rrrpr4rgrYrZr)rbtuprUrVr\rjrCrmmanifest_entryraudio_filestreamrrrrqrrr&r&r*rHs8       " z'_TarredAudioToTextDataset._build_samplecCrrfrrr&r&r*rrz-_TarredAudioToTextDataset.get_manifest_samplecCs |jSrf)rIrPrr&r&r*rPs z"_TarredAudioToTextDataset.__iter__cCsn|jr/tjr/tjr/tjt|jjtj d }tj || }t d||St|jj}|S)Nr0z!Sharded manifests: Total length: )rr8rrrr>r4rrar?cuda all_reduceryrr)rbmy_lenr&r&r*r<s  z&_TarredAudioToTextDataset._compute_lencCs|jSrf)r4rr&r&r*rsz!_TarredAudioToTextDataset.__len__)FNrNNFNNrrFrrFN)rsrtrurvrrwrrryrzrrxrcrErFrrHrrPr<rr&r&r&r*r5sb      N- r5c0seZdZdZ                   d#d eeeefd ed eed ed ede ddede e de e dedededede ede ede ededededededed e e f.fd!d" Z Z S)$ra A similar Dataset to the AudioToCharDataset, but which loads tarred audio files. Accepts a single comma-separated JSON manifest file (in the same style as for the AudioToCharDataset), as well as the path(s) to the tarball(s) containing the wav files. Each line of the manifest should contain the information for one audio file, including at least the transcript and name of the audio file within the tarball. Valid formats for the audio_tar_filepaths argument include: (1) a single string that can be brace-expanded, e.g. 'path/to/audio.tar' or 'path/to/audio_{1..100}.tar.gz', or (2) a list of file paths that will not be brace-expanded, e.g. ['audio_1.tar', 'audio_2.tar', ...]. See the WebDataset documentation for more information about accepted data and input formats. If using multiple workers the number of shards should be divisible by world_size to ensure an even split among workers. If it is not divisible, logging will give a warning but training will proceed. In addition, if using mutiprocessing, each shard MUST HAVE THE SAME NUMBER OF ENTRIES after filtering is applied. We currently do not check for this, but your program may hang if the shards are uneven! Notice that a few arguments are different from the AudioToCharDataset; for example, shuffle (bool) has been replaced by shuffle_n (int). Additionally, please note that the len() of this DataLayer is assumed to be the length of the manifest after filtering. An incorrect manifest length may lead to some DataLoader issues down the line. Args: audio_tar_filepaths: Either a list of audio tarball filepaths, or a string (can be brace-expandable). manifest_filepath (str): Path to the manifest. labels (list): List of characters that can be output by the ASR model. For Jasper, this is the 28 character set {a-z '}. The CTC blank symbol is automatically added later for models using ctc. sample_rate (int): Sample rate to resample loaded audio to int_values (bool): If true, load samples as 32-bit integers. Defauts to False. augmentor (nemo.collections.asr.parts.perturb.AudioAugmentor): An AudioAugmentor object used to augment loaded audio shuffle_n (int): How many samples to look ahead and load to be shuffled. See WebDataset documentation for more details. Defaults to 0. min_duration (float): Dataset parameter. All training files which have a duration less than min_duration are dropped. Note: Duration is read from the manifest JSON. Defaults to 0.1. max_duration (float): Dataset parameter. All training files which have a duration more than max_duration are dropped. Note: Duration is read from the manifest JSON. Defaults to None. blank_index (int): Blank character index, defaults to -1. unk_index (int): Unknown character index, defaults to -1. normalize (bool): Dataset parameter. Whether to use automatic text cleaning. It is highly recommended to manually clean text for best results. Defaults to True. trim (bool): Whether to use trim silence from beginning and end of audio signal using librosa.effects.trim(). Defaults to False. bos_id (id): Dataset parameter. Beginning of string symbol id used for seq2seq models. Defaults to None. eos_id (id): Dataset parameter. End of string symbol id used for seq2seq models. Defaults to None. pad_id (id): Token used to pad when collating samples in batches. If this is None, pads using 0s. Defaults to None. shard_strategy (str): Tarred dataset shard distribution strategy chosen as a str value during ddp. - `scatter`: The default shard strategy applied by WebDataset, where each node gets a unique set of shards, which are permanently pre-allocated and never changed at runtime. - `replicate`: Optional shard strategy, where each node gets all of the set of shards available in the tarred dataset, which are permanently pre-allocated and never changed at runtime. The benefit of replication is that it allows each node to sample data points from the entire dataset independently of other nodes, and reduces dependence on value of `shuffle_n`. .. warning:: Replicated strategy allows every node to sample the entire set of available tarfiles, and therefore more than one node may sample the same tarfile, and even sample the same data points! As such, there is no assured guarantee that all samples in the dataset will be sampled at least once during 1 epoch. Scattered strategy, on the other hand, on specific occasions (when the number of shards is not divisible with ``world_size``), will not sample the entire dataset. For these reasons it is not advisable to use tarred datasets as validation or test datasets. global_rank (int): Worker rank, used for partitioning shards. Defaults to 0. world_size (int): Total number of processes, used for partitioning shards. Defaults to 0. return_sample_id (bool): whether to return the sample_id as a part of each sample manifest_parse_func: Optional function to parse manifest entries. Defaults to None. FNrrTrrr6rTrrrrrr7rWrVr r r rrYrZrUrAr|rr~r}rr\cs||_tj||| | | d}tjdid|d|d|d|d|d|d|d |d | d | d |d |d|d|d|d|d|d|d|dS)Nr r6rTrUrrrr7rWrVrrYrZrAr|rr~r}rr\r&r)rbr6rTrrrrr7rWrVr r r rrYrZrUrAr|rr~r}rr\rr&r*rcVsV       z!TarredAudioToCharDataset.__init__)FNrNNrrTFNNrrrFrrFNrsrtrurvrrwrryrzrrxrrcrr&r&rr*rs`     rc$seZdZdZ             ddeeeefded d d ed ed e ddede e de e dededededededede e f"fdd Z Z S)r aS A similar Dataset to the AudioToBPEDataset, but which loads tarred audio files. Accepts a single comma-separated JSON manifest file (in the same style as for the AudioToBPEDataset), as well as the path(s) to the tarball(s) containing the wav files. Each line of the manifest should contain the information for one audio file, including at least the transcript and name of the audio file within the tarball. Valid formats for the audio_tar_filepaths argument include: (1) a single string that can be brace-expanded, e.g. 'path/to/audio.tar' or 'path/to/audio_{1..100}.tar.gz', or (2) a list of file paths that will not be brace-expanded, e.g. ['audio_1.tar', 'audio_2.tar', ...]. See the WebDataset documentation for more information about accepted data and input formats. If using multiple workers the number of shards should be divisible by world_size to ensure an even split among workers. If it is not divisible, logging will give a warning but training will proceed. In addition, if using mutiprocessing, each shard MUST HAVE THE SAME NUMBER OF ENTRIES after filtering is applied. We currently do not check for this, but your program may hang if the shards are uneven! Notice that a few arguments are different from the AudioToBPEDataset; for example, shuffle (bool) has been replaced by shuffle_n (int). Additionally, please note that the len() of this DataLayer is assumed to be the length of the manifest after filtering. An incorrect manifest length may lead to some DataLoader issues down the line. Args: audio_tar_filepaths: Either a list of audio tarball filepaths, or a string (can be brace-expandable). manifest_filepath (str): Path to the manifest. tokenizer (TokenizerSpec): Either a Word Piece Encoding tokenizer (BERT), or a Sentence Piece Encoding tokenizer (BPE). The CTC blank symbol is automatically added later for models using ctc. sample_rate (int): Sample rate to resample loaded audio to int_values (bool): If true, load samples as 32-bit integers. Defauts to False. augmentor (nemo.collections.asr.parts.perturb.AudioAugmentor): An AudioAugmentor object used to augment loaded audio shuffle_n (int): How many samples to look ahead and load to be shuffled. See WebDataset documentation for more details. Defaults to 0. min_duration (float): Dataset parameter. All training files which have a duration less than min_duration are dropped. Note: Duration is read from the manifest JSON. Defaults to 0.1. max_duration (float): Dataset parameter. All training files which have a duration more than max_duration are dropped. Note: Duration is read from the manifest JSON. Defaults to None. trim (bool): Whether to use trim silence from beginning and end of audio signal using librosa.effects.trim(). Defaults to False. use_start_end_token: Boolean which dictates whether to add [BOS] and [EOS] tokens to beginning and ending of speech respectively. pad_id (id): Token used to pad when collating samples in batches. If this is None, pads using 0s. Defaults to None. shard_strategy (str): Tarred dataset shard distribution strategy chosen as a str value during ddp. - `scatter`: The default shard strategy applied by WebDataset, where each node gets a unique set of shards, which are permanently pre-allocated and never changed at runtime. - `replicate`: Optional shard strategy, where each node gets all of the set of shards available in the tarred dataset, which are permanently pre-allocated and never changed at runtime. The benefit of replication is that it allows each node to sample data points from the entire dataset independently of other nodes, and reduces dependence on value of `shuffle_n`. .. warning:: Replicated strategy allows every node to sample the entire set of available tarfiles, and therefore more than one node may sample the same tarfile, and even sample the same data points! As such, there is no assured guarantee that all samples in the dataset will be sampled at least once during 1 epoch. Scattered strategy, on the other hand, on specific occasions (when the number of shards is not divisible with ``world_size``), will not sample the entire dataset. For these reasons it is not advisable to use tarred datasets as validation or test datasets. global_rank (int): Worker rank, used for partitioning shards. Defaults to 0. world_size (int): Total number of processes, used for partitioning shards. Defaults to 0. return_sample_id (bool): whether to return the sample_id as a part of each sample manifest_parse_func: Optional function to parse manifest entries. Defaults to None. FNrTrr6rTrrrrrrr7rWrVrrr|rr~r}rr\cs| rt|dr|jdkr|j}nd}| r"t|dr"|jdkr"|j}nd}t|dr2|jdkr2|j}nd}Gddd}tjdid|d|d ||d |d |d |d |d|d| d| d|d|d|d| d| d|d|d|d|dS)NrYrrZrAc@r)z:TarredAudioToBPEDataset.__init__..TokenizerWrappercSrrrr"r&r&r*rcr#zCTarredAudioToBPEDataset.__init__..TokenizerWrapper.__init__cWr$r%r'r*r&r&r*r- r.zCTarredAudioToBPEDataset.__init__..TokenizerWrapper.__call__Nr/r&r&r&r*r0r1r0r6rTrUrrrr7rWrVrr|rr~r}rr\r&r2)rbr6rTrrrrr7rWrVrrr|rr~r}rr\rYrZrAr0rr&r*rcsb      z TarredAudioToBPEDataset.__init__) FNrNNFTrFrrFNrhr&r&rr*r sdV     r csBeZdZdZdedeffdd ZddZdd Zd d Z Z S) BucketingDatasetz A Dataset which wraps another IterableDataset and adopts it for bucketing Args: dataset (IterableDataset): The IterableDataset to get wrapped bucketing_batch_size (int): Number of samples to build a batch datasetbucketing_batch_sizecs||_||_tdSrf)wrapped_datasetrkrrc)rbrjrkrr&r*rc2szBucketingDataset.__init__cCst|d|jjSrk)rRrlrArr&r&r*r;szBucketingDataset._collate_fncCst|jj|jdS)N) wrapped_dsrk)BucketingIteratorrlrIrkrPrr&r&r*rP>s  zBucketingDataset.__iter__cCsttt|jt|jSrf)rymathceilr4rlrxrkrr&r&r*rCszBucketingDataset.__len__) rsrtrurvrryrcrrPrrr&r&rr*ri*s ric@s$eZdZddZddZddZdS)rncCs||_d|_||_dSrf)rm wrapped_iterrk)rbrmrkr&r&r*rcHs zBucketingIterator.__init__cCst|j|_|Srf)iterrmrqrr&r&r*rPMs zBucketingIterator.__iter__c CsVg}t|jD]}zt|j}Wn tyYnw||qt|dkr)t|Srk)rangerkrRrq StopIterationr<r4)rbbatchesidxrhr&r&r*rWQs   zBucketingIterator.__next__NrXr&r&r&r*rnGs rncs4eZdZd deeddffdd ZddZZS) RandomizedChainDatasetrdatasetsreNcs&tt|t|tj||_dSrf)rrwrcr2nprandom RandomStaternd_gen)rbrxrnd_seedrr&r*rc_szRandomizedChainDataset.__init__ccsj|jt|j}|D]&}|j|}t|tsJdt|D]\}}|V|t|dkr1nq q dS)Nz*ChainDataset only supports IterableDatasetro)r| permutationr4rxrr enumerate)rbshuffled_order dataset_idxdrvxr&r&r*rPcs zRandomizedChainDataset.__iter__)r)rsrtrurrrcrPrr&r&rr*rw^srw)FNNr)Ir_rrorrcollections.abcrrtypingrrrrrrrnumpyryr8torch.utils.datar r 1nemo.collections.asr.parts.preprocessing.featuresr 0nemo.collections.asr.parts.preprocessing.segmentr r valid_sf_formatsnemo.collections.commonr+nemo.collections.common.parts.preprocessingrrnemo.core.classesrrnemo.core.neural_types nemo.utilsrrr=nemo.utils.data_utilsrrrrnemo.utils.decoratorsrnemo.utils.distributedrnemo.utils.get_rankr__all__rkeysrCrRrSrwryrrzrrrrrr5rr rirnrwr&r&r&r*s  $            =F2 x  x\u)