o }oig @sddlZddlZddlZddlmZmZddlmZddlm Z ddl m Z m Z m Z mZddlZddlmZddlmZmZmZmZddlmZdd lmZdd lmZdd lmZmZdd l m!Z!dd l"m#Z#m$Z$ddl%m&Z&ddl'm(Z(ddl)m*Z*GdddZ+GdddZ, d,dede-de-dBdefddZ.GdddZ/Gddde/Z0Gd d!d!e/Z1d"d#Z2d$e3e Be4e3Bde4e3fd%d&Z5dd'hfd(e6d)e7e3de6fd*d+Z8dS)-N)MappingSequence)BytesIO)Path) GeneratorIterableListLiteral)groupby) AudioSourceMonoCut RecordingSupervisionSegment)LibsndfileBackend)Cut) resolve_seed)LazyIteratorChainLazyJsonlIterator) open_best)compute_num_samplesifnone) get_full_path)logging)is_datastore_pathc@seZdZdZ      d deeBeeBded ed ed ed ee d Bdee eefdBddfddZ de e ddffddZdefddZddZ d!dededededBde f ddZ d!dedededBdefddZdS)"LazyNeMoIteratora# ``LazyNeMoIterator`` reads a NeMo (non-tarred) JSON manifest and converts it on the fly to an ``Iterable[Cut]``. It's used to create a ``lhotse.CutSet``. Currently, it requires the following keys in NeMo manifests: - "audio_filepath" - "duration" - "text" (overridable with ``text_field`` argument) Specially supported keys are: - [recommended] "sampling_rate" allows us to provide a valid Lhotse ``Recording`` object without checking the audio file - "offset" for partial recording reads - "lang" is mapped to Lhotse superivsion's language (overridable with ``lang_field`` argument) Every other key found in the manifest will be attached to Lhotse Cut and accessible via ``cut.custom[key]``. .. caution:: We will perform some I/O (as much as required by soundfile.info) to discover the sampling rate of the audio file. If this is not acceptable, convert the manifest to Lhotse format which contains sampling rate info. For pure metadata iteration purposes we also provide a ``metadata_only`` flag that will create only partially valid Lhotse objects (with metadata related to sampling rate / num samples missing). Example:: >>> cuts = lhotse.CutSet(LazyNeMoIterator("nemo_manifests/train.json")) We allow attaching custom metadata to cuts from files other than the manifest via ``extra_fields`` argument. In the example below, we'll iterate file "questions.txt" together with the manifest and attach each line under ``cut.question`` using the field type ``text_iter``:: >>> cuts = lhotse.CutSet(LazyNeMoIterator( ... "nemo_manifests/train.json", ... extra_fields=[{"type": "text_iter", "name": "question", "path": "questions.txt"}], ... )) We also support random sampling of lines with field type ``text_sample``:: >>> cuts = lhotse.CutSet(LazyNeMoIterator( ... "nemo_manifests/train.json", ... extra_fields=[{"type": "text_sample", "name": "question", "path": "questions.txt"}], ... )) textlangFtrngNpath text_field lang_field metadata_onlyshuffle_shards shard_seed) randomizedr extra_fieldsreturnc Cs|||_||_||_t|}t|dkrt|d|_ntdd|D|j|jd|_||_||_ ||_ ||_ t |j dS)Nrcss|]}t|VqdSNr).0pr,e/home/ubuntu/.local/lib/python3.10/site-packages/nemo/collections/common/data/lhotse/nemo_adapters.py dsz,LazyNeMoIterator.__init__..) shuffle_itersseed) rr"r#expand_sharded_filepathslenrsourcerrr r!r%validate_extra_fields) selfrrr r!r"r#r%pathsr,r,r-__init__Qs   zLazyNeMoIterator.__init__c #st|jfdd|jpdD}|jD]Y}|ddrqtt|dt|jdd}|d}|d d}|j ||||d dd }|j t |j |jd |j||j||jd ||_|D]}||qc|VqdS)Ncg|] }tdi|qSr0 ExtraField from_dictr* field_cfgr9r,r- oz-LazyNeMoIterator.__iter__..r,_skipmeFaudio_filepath) force_cachedurationoffset sampling_rate) audio_pathrErDrFrid recording_idstartrDrlanguage)rr#r%r3getrstrpopr _create_cut supervisionsappendrrIrJrDrr custom attach_to)r5r%datarGrDrEcut extra_fieldr,r9r-__iter__ls6         zLazyNeMoIterator.__iter__cC t|jSr(r2r3r5r,r,r-__len__ zLazyNeMoIterator.__len__cC t||Sr(rr5otherr,r,r-__add__r]zLazyNeMoIterator.__add__rGrErDrFcCs|js2||||}|}|dur0|j||dd}|jdt|dddt|dd|_|St|d}t|d}t|||dgt|t d dgd d g|||t |||d d }|S)NTrErD preserve_id-gY@06di>rdummytypechannelsr3)rIsourcesrFrD num_samples)rIrKrDchannelrQ recording) r!_create_recordingto_cuttruncaterIroundrr r r r)r5rGrErDrFrprVsrr,r,r-rPs0,  zLazyNeMoIterator._create_cutcCsL|dur!t|r dnd}t|t|dg|dg|t|||dgdSt|S)Nurlfilerrj)rIrmrFrnrD channel_ids)rr r r from_file)r5rGrDrF source_typer,r,r-rqs z"LazyNeMoIterator._create_recording)rrFFrNr()__name__ __module__ __qualname____doc__rNrlistboolintr dictr7rrrXr\rbfloatrPr rqr,r,r,r-r&sd-   (rc@seZdZdZ      d!deeBeeBdeeBd ed ee d Bd ed ededee eefdBddfddZ de dfddZ d"ddZede efddZdeee efddffddZdeeddffddZdefddZdd ZdS)#LazyNeMoTarredIteratora ``LazyNeMoTarredIterator`` reads a NeMo tarred JSON manifest and converts it on the fly to an ``Iterable[Cut]``. It's used to create a ``lhotse.CutSet``. Currently, it requires the following keys in NeMo manifests: - "audio_filepath" - "duration" - "text" (overridable with text_field argument) - "shard_id" Specially supported keys are: - "lang" is mapped to Lhotse superivsion's language (overridable with ``lang_field`` argument) Every other key found in the manifest will be attached to Lhotse Cut and accessible via ``cut.custom[key]``. Args ``manifest_path`` and ``tar_paths`` can be either a path/string to a single file, or a string in NeMo format that indicates multiple paths (e.g. "[[data/bucket0/tarred_audio_paths.json],[data/bucket1/...]]"). We discover shard ids from sharded tar and json files by parsing the input specifier/path and searching for the following pattern: ``(manifest|audio)[^/]*_(\d+)[^/]*\.(json|tar)``. It allows filenames such as ``manifest_0.json``, ``manifest_0_normalized.json``, ``manifest_normalized_0.json``, ``manifest_0.jsonl.gz``, etc. (anologusly the same applies to tar files). We also support generalized input specifiers that imitate webdataset's pipes (also very similar to Kaldi's pipes). These are arbitrary shell commands to be lazily executed which yield manifest or tar audio contents. For example, ``tar_paths`` can be set to ``pipe:ais get ais://my-bucket/audio_{0..127}.tar -`` to indicate that we want to read tarred audio data from shards on an AIStore bucket. This can be used for other cloud storage APIs such as S3, GCS, etc. The same mechanism applies to ``manifest_path``. If your data has been filtered so that the JSON manifests refer to just a subset of recordings, set ``skip_missing_manifest_entries` to ``True``. This will still read the tar files sequentially (very fast) and discard the audio files that are not present in the corresponding manifest. The ``shard_seed`` argument is used to seed the RNG shuffling the shards. By default, it's ``trng`` which samples a seed number from OS-provided TRNG (see Python ``secrets`` module). Seed is resolved lazily so that every dataloading worker may sample a different one. Override with an integer value for deterministic behaviour and consult Lhotse documentation for details: https://lhotse.readthedocs.io/en/latest/datasets.html#handling-random-seeds Example of CutSet with inter-shard shuffling enabled:: >>> cuts = lhotse.CutSet(LazyNeMoTarredIterator( ... manifest_path=["nemo_manifests/sharded_manifests/manifest_0.json", ...], ... tar_paths=["nemo_manifests/audio_0.tar", ...], ... shuffle_shards=True, ... )) We allow attaching custom metadata to cuts from files other than the manifest via ``extra_fields`` argument. In the example below, we'll iterate file "questions.txt" together with the manifest and attach each line under ``cut.question`` using the field type ``text_iter``:: >>> cuts = lhotse.CutSet(LazyNeMoTarredIterator( ... manifest_path=["nemo_manifests/sharded_manifests/manifest_0.json", ...], ... tar_paths=["nemo_manifests/audio_0.tar", ...], ... extra_fields=[{"type": "text_iter", "name": "question", "path": "questions.txt"}], ... )) We also support random sampling of lines with field type ``text_sample``:: >>> cuts = lhotse.CutSet(LazyNeMoTarredIterator( ... manifest_path=["nemo_manifests/sharded_manifests/manifest_0.json", ...], ... tar_paths=["nemo_manifests/audio_0.tar", ...], ... extra_fields=[{"type": "text_sample", "name": "question", "path": "questions.txt"}], ... )) FrrrN manifest_path tar_pathsr"r#)rr$rr skip_missing_manifest_entriesr%r&c Csx||_|t||_t|jdkr-td|jddt|jd|_td|j|_ n@t d} g} |jD]!} | | } | dusNJd| j d| d | t| dq7d d t| |jD|_ t|j |_t||_t d } g} |jD]!} | | } | dusJd | j d| d | t| dq|tt| |j|_||_||_||_||_||_|dS)Nr'zYou are using Lhotse dataloading for tarred audio with a non-sharded manifest. This will incur significant memory overhead and slow-down training. To prevent this error message please shard file 'rz' using 'scripts/speech_recognition/convert_to_tarred_audio_dataset.py' WITHOUT '--no_shard_manifest'shard_idzmanifest[^/]*_(\d+)[^/]*\.jsonzQCannot determine shard_id from manifest input specified: we searched with regex 'z ' in input ''cSsi|] \}}|t|qSr,r))r*sidr+r,r,r- 0sz3LazyNeMoTarredIterator.__init__..zaudio[^/]*_(\d+)[^/]*\.tarzLCannot determine shard_id from tar input specifier: we searched with regex ')rr1r6r2rwarningrr3r shard_id_to_manifestrecompilesearchpatternrRrgroupziprvaluesrrshard_id_to_tar_pathr"r#rr r% _validate)r5rrr"r#rr rr% json_pattern shard_idsr+m tar_patternr,r,r-r7sZ            zLazyNeMoTarredIterator.__init__cs4tjdkr gSfddtjjDS)zEConvert this iterator to a list of separate iterators for each shard.r'c s*g|]\}}t||djjjdqS)F)rrr"r#rr )rr#rr )r*rtarpathr[r,r-r?Ms z4LazyNeMoTarredIterator.to_shards..)r2r6rrrr[r,r[r- to_shardsFs   z LazyNeMoTarredIterator.to_shardsc CsXt|j}t|j}||ks%Jd|jd|jdt|dt|d t|jdS)Nz8Mismatch between shard IDs. Details: * JSON manifest(s) z * Tar files: z% * JSON manifest(s) indicate(s) IDs: z * Tar path(s) indicate(s) IDs:  )setrrr6rsortedr4r%)r5shard_ids_tarsshard_ids_manifestr,r,r-rYs   z LazyNeMoTarredIterator._validatecCst|jSr()rrkeysr[r,r,r-resz LazyNeMoTarredIterator.shard_idsc cstjt|ddddG}|D];}z||j}||}|||fVWqtyJ}z|jr6WYd}~qtd|d|d|jd|d}~wwWddS1sVwYdS) Nrb)modezr|*)fileobjrz)Mismatched entry between JSON manifest ('z') and tar file ('z+'). Cannot locate JSON entry for tar file 'r) tarfileopenrname extractfilereadKeyErrorr RuntimeError) r5tar_pathshard_manifestrtartar_inforU raw_audioer,r,r-_iter_sequentialis*  "z'LazyNeMoTarredIterator._iter_sequentialc#s|j}t|j|jrt|fdd|jpdD}t d|D]}t |j dkr5|j |n|j d}dt dt ffd d }t||j|}|j|}z||||D]\}} } tt| } t| jtd tt| j| d gt| j| j| jd } g} t|| j dddD]J}|!ddrqt"| |!dd|!dd}|j#$t%|j&|j'd|j|!|j(|!|j)dt*||_+||_,||_-|D]}|.|q| $|q~ ~ | EdHqZWq't/j0yt12d|Yq'wdS)Ncr8r9r:r=r9r,r-r?r@z3LazyNeMoTarredIterator.__iter__..r,z-^(?P.+)(?P-sub\d+)(?P\.\w+)?$r'rdr&cs8|d}}dur|dt|ddS|S)NrBstemextri)matchrr)rkr)offset_patternr,r-basenames z1LazyNeMoTarredIterator.__iter__..basenamememoryrj)rIrmrFrnrDcSs|dS)NrBr,)rr,r,r-sz1LazyNeMoTarredIterator.__iter__..)keyrAFrErgrD)rErDrHzOSkipping tar file due to read errors (unstable storage or bad file?): tar_path=)3rrr#r"randomRandomshuffler%rrr2r6rrNr rrr soundfileinforr rr rrangerlr samplerateframesrDrrrM'make_cut_with_subset_inmemory_recordingrQrRrrIrJrr _to_custom_attr_dictrSmanifest_origin tar_originrTr ReadErrorrr)r5rr%rrrrrrUrrmetarpcuts_for_recordingrVrWr,)rr0r-rXysn  "        %zLazyNeMoTarredIterator.__iter__cCrYr(rZr[r,r,r-r\r]zLazyNeMoTarredIterator.__len__cCr^r(r_r`r,r,r-rbr]zLazyNeMoTarredIterator.__add__)FrrrFN)r&N)r{r|r}r~rNrrrrr rr7rrrpropertyrrtuplebytesrrrXr\rbr,r,r,r-rsFG   7   E rrgrprErDr&c Cs|}|dkr |dust||jdkr|Sz |j||dd}Wnty>}ztd|d|d|d ||d}~wwt}tj|| |j d d | d t |j |j |j|jtd |j|dgd}|S)a This method is built specifically to optimize CPU memory usage during dataloading when reading tarfiles containing very long recordings (1h+). Normally each cut would hold a reference to the long in-memory recording and load the necessary subset of audio (there wouldn't be a separate copy of the long recording for each cut). This is fairly efficient already, but we don't actually need to hold the unused full recording in memory. Instead, we re-create each cut so that it only holds a reference to the subset of recording necessary. This allows us to discard unused data which would otherwise be held in memory as part of sampling buffering. rgNg?Trcz'Lhotse cut.truncate failed with offset=z , duration=z , recording=z: wav)rFformatrrrj)rIrFrnrDrm)rrabsrDrs Exceptionrrr save_audio load_audiorFseekr rIrnr rxgetvalue)rprErDrVr audiobytes new_recordingr,r,r-rs:"  rcsneZdZdZiZddZfddZededdfdd Z e d e de fd d Z e dee fd dZZS)r;NcCstr()NotImplementedError)r5rVr,r,r-rTszExtraField.attach_toc s.|jtjvr |tj|j<tjdi|dS)Nr,)r{r;SUPPORTED_TYPESTYPEsuper__init_subclass__)clskwargs __class__r,r-rs  zExtraField.__init_subclass__rUr&cCsF|dtjvsJd|dtj|ddidd|DS)NrkzUnknown transform type: cSsi|] \}}|dkr||qS)rkr,r*rvr,r,r-rr@z(ExtraField.from_dict..r,)r;ritems)rUr,r,r-r<s &zExtraField.from_dict field_typecCs ||jvSr()r)rrr,r,r- is_supported zExtraField.is_supportedcCrYr()rr)rr,r,r-supported_typesrzExtraField.supported_types)r{r|r}rrrTr staticmethodrr< classmethodrNrrrr __classcell__r,r,rr-r;s r;c@s4eZdZdZd dedefddZddZd d ZdS) TextIteratorExtraField text_iterNrrcCs||_||_d|_dSr()rriteratorr5rrr0r,r,r-r7s zTextIteratorExtraField.__init__cCs*|jdurtttjt|j|_dSdSr()ritermaprNstriprrr[r,r,r- _maybe_inits z"TextIteratorExtraField._maybe_initc CsR|zt|j}Wntytd|jd|jdwt||j||S)NzNot enough lines in file z to attach to cuts under field .)rnextr StopIterationrrrsetattrr5rVattached_valuer,r,r-rTs z TextIteratorExtraField.attach_tor()r{r|r}rrNr7rrTr,r,r,r-r s  rc@s:eZdZdZdededeeBfddZddZd d Zd S) TextSampleExtraField text_samplerrr0cCs"||_||_||_d|_d|_dSr()rrr0 populationrngrr,r,r-r7%s  zTextSampleExtraField.__init__cCs<|jdurtttjt|j|_tt |j |_ dSdSr() rrrrNrrrrrrr0rr[r,r,r-r,s z TextSampleExtraField._maybe_initcCs(||j|j}t||j||Sr()rrchoicerrrrr,r,r-rT1szTextSampleExtraField.attach_toN) r{r|r}rrNrr7rrTr,r,r,r-r"s  rcCs|durdSt|tsJd||D]8}t|ts%Jd|d||d}t|s>Jdtd|d|d|vsLJd|d|qdS) NzZThe argument provided to 'extra_fields' must be a list of dicts. We received extra_fields=z>Each item in 'extra_fields' must be a dict. We received field=z in extra_fields=rkzZEach item in 'extra_fields' must contain a 'type' field with one of the supported values (z). We got field_type=rzwEach item in 'extra_fields' must contain a 'name' field so that the field is available under cut..We found field=) isinstancerrrMr;rr)r%fieldrr,r,r-r48s>   r4r6cCs.ddlm}t|trt|}||ddddS)Nr)r1 replicater')shard_strategy world_size global_rank)'nemo.collections.asr.data.audio_to_textr1rrrN)r6_expand_sharded_filepathsr,r,r-r1Ns  r1rBr_excluded_fieldscsfdd|DS)Ncsi|] \}}|vr||qSr,r,rrr,r-rYr@z(_to_custom_attr_dict..)r)rrr,r r-rXsr)rgN)9rrrcollections.abcrriorpathlibrtypingrrrr rcytoolzr lhotser r r rlhotse.audio.backendr lhotse.cutrlhotse.dataset.dataloadingr lhotse.lazyrrlhotse.serializationr lhotse.utilsrr4nemo.collections.common.parts.preprocessing.manifestr nemo.utilsrnemo.utils.data_utilsrrrrrr;rrr4rNrr1rrrr,r,r,r-sN          &| -"&