o Si@sddlZddlmZmZddlmZddlmZddlm Z ddl m Z m Z ddl mZddlmZmZmZmZmZddlZddlZdd lmZmZdd lmZdd lmZdd lm Z m!Z!m"Z"m#Z#m$Z$m%Z%Gd ddedZ&GdddedZ'iZ(iZ)dee*fddZ+ddZ,ddZ-de*dee'fddZ.de*dee&fddZ/GdddZ0 e,Gd d!d!e'Z1e-Gd"d#d#e&Z2 e,Gd$d%d%e'Z3e-Gd&d'd'e&Z4 d(d)Z5edd*d+e*fd,d-Z6edd*de7fd.d/Z8dad0d1Z9e,Gd2d3d3e'Z:e-Gd4d5d5e&Z; e,Gd6d7d7e'Ze,Gd;d<dd>e&Z@ d?ZAe,Gd@dAdAe'ZBe-GdBdCdCe&ZC e,GdDdEdEe'ZDe-GdFdGdGe&ZE dHdIZFedd*d+e*fdJdKZGe,GdLdMdMe'ZHe-GdNdOdOe&ZI dPe*deJfdQdRZKde*fdSdTZLe,GdUdVdVe'ZMe-GdWdXdXe&ZNe,GdYdZdZe'ZOe-Gd[d\d\e&ZPe,Gd]d^d^e'ZQe,Gd_d`d`e'ZRdS)bN)ABCMetaabstractmethod)contextmanager) lru_cache)BytesIO)ceilfloor)Path) GeneratorListOptionalTypeUnion)Array TemporalArray)dynamic_lru_cache) open_best)PathlikeSeconds SmartOpenis_module_available is_valid_urlpairwisec@seZdZdZeedefddZeedefddZedede j defd d Z ddede j d e e de ede deeeff ddZddZddZd S)FeaturesWritera} ``FeaturesWriter`` defines the interface of how to store numpy arrays in a particular storage backend. This backend could either be: - separate files on a local filesystem; - a single file with multiple arrays; - cloud storage; - etc. Each class inheriting from ``FeaturesWriter`` must define: - the ``write()`` method, which defines the storing operation (accepts a ``key`` used to place the ``value`` array in the storage); - the ``storage_path()`` property, which is either a common directory for the files, the name of the file storing multiple arrays, name of the cloud bucket, etc. - the ``name()`` property that is unique to this particular storage mechanism - it is stored in the features manifests (metadata) and used to automatically deduce the backend when loading the features. Each :class:`.FeaturesWriter` can also be used as a context manager, as some implementations might need to free a resource after the writing is finalized. By default nothing happens in the context manager functions, and this can be modified by the inheriting subclasses. Example:: >>> with MyWriter('some/path') as storage: ... extractor.extract_from_recording_and_store(recording, storage) The features loading must be defined separately in a class inheriting from :class:`FeaturesReader`. returncCdSNselfrrF/home/ubuntu/.local/lib/python3.10/site-packages/lhotse/features/io.pyname:zFeaturesWriter.namecCrrrrrrr storage_path?r"zFeaturesWriter.storage_pathkeyvaluecCrrrrr$r%rrr writeDzFeaturesWriter.writeNr frame_shift temporal_dimstartc Cs~|duo|du}|s tdd||fDs Jd|d|d|||}t|j|j|t|jd}|s7|St||||dS) aD Store a numpy array in the underlying storage and return a manifest describing how to retrieve the data. If the array contains a temporal dimension (e.g. it represents the frame-level features, alignment, posteriors, etc. of an utterance) then ``temporal_dim`` and ``frame_shift`` may be specified to enable downstream padding, truncating, and partial reads of the array. :param key: An ID that uniquely identifies the array. :param value: The array to be stored. :param frame_shift: Optional float, when the array has a temporal dimension it indicates how much time has passed between the starts of consecutive frames (expressed in seconds). :param temporal_dim: Optional int, when the array has a temporal dimension, it indicates which dim to interpret as temporal. :param start: Float, when the array is temporal, it indicates what is the offset of the array w.r.t. the start of recording. Useful for reading subsets of an array when it represents something computed from long recordings. Ignored for non-temporal arrays. :return: A manifest of type :class:`~lhotse.array.Array` or :class:`~lhotse.array.TemporalArray`, depending on the input arguments. Ncss|]}|duVqdSrr).0argrrr isz-FeaturesWriter.store_array..zOframe_shift and temporal_dim have to be both None or both set (got frame_shift=z, temporal_dim=z).) storage_typer# storage_keyshape)arrayr*r)r+)allr'rr!r#listr1r) rr$r%r)r*r+ is_temporalr0r2rrr store_arrayHs0 zFeaturesWriter.store_arraycC|Srrrrrr __enter__zFeaturesWriter.__enter__cOrrrrargskwargsrrr __exit__r9zFeaturesWriter.__exit__)NNr)__name__ __module__ __qualname____doc__propertyrstrr!r#npndarrayr'r rintrrrr6r8r=rrrr rs8  8 r) metaclassc @sPeZdZdZeedefddZe  d deded e ede j fd d Z dS) FeaturesReadera ``FeaturesReader`` defines the interface of how to load numpy arrays from a particular storage backend. This backend could either be: - separate files on a local filesystem; - a single file with multiple arrays; - cloud storage; - etc. Each class inheriting from ``FeaturesReader`` must define: - the ``read()`` method, which defines the loading operation (accepts the ``key`` to locate the array in the storage and return it). The read method should support selecting only a subset of the feature matrix, with the bounds expressed as arguments ``left_offset_frames`` and ``right_offset_frames``. It's up to the Reader implementation to load only the required part or trim it to that range only after loading. It is assumed that the time dimension is always the first one. - the ``name()`` property that is unique to this particular storage mechanism - it is stored in the features manifests (metadata) and used to automatically deduce the backend when loading the features. The features writing must be defined separately in a class inheriting from ``FeaturesWriter``. rcCrrrrrrr r!r"zFeaturesReader.namerNr$left_offset_framesright_offset_framescCrrrrr$rIrJrrr readszFeaturesReader.readrN) r>r?r@rArBrrCr!rFr rDrErLrrrr rHs"rHrcCsttttSr)sortedsetREADER_BACKENDS intersectionWRITER_BACKENDSrrrr available_storage_backendssrScC|t|j<|S)z Decorator used to add a new ``FeaturesReader`` to Lhotse's registry. Example:: @register_reader class MyFeatureReader(FeatureReader): ... )rPr!clsrrr register_reader rWcCrT)z Decorator used to add a new ``FeaturesWriter`` to Lhotse's registry. Example:: @register_writer class MyFeatureWriter(FeatureWriter): ... )rRr!rUrrr register_writerrXrYr!cC t|S)z Find a ``FeaturesReader`` sub-class that corresponds to the provided ``name`` and return its type. Example: reader_type = get_reader("lilcom_files") reader = reader_type("/storage/features/") )rPgetr!rrr get_reader r]cCrZ)z Find a ``FeaturesWriter`` sub-class that corresponds to the provided ``name`` and return its type. Example: writer_type = get_writer("lilcom_files") writer = writer_type("/storage/features/") )rRr[r\rrr get_writerr^r_c sReZdZdZdeffdd Ze d dededed e e d d ffd d Z Z S)FileIOaG Helper util for opening a file object for reading or writing in a directory on the local filesystem, or a URL to supported object store (S3, AIStore, etc.). ``storage_path`` corresponds to the directory path or base URL prefix; ``storage_key`` for each utterance is the name of the file in that directory. r#csLtt||_t||_|jr"|jdr$|jdd|_dSdSdS)N/)super__init__rCr#ris_urlendswith)rr# __class__rr rds    zFileIO.__init__Fr$mode add_subdirrNc csZd|vr d|vr Jdd|vrH|dr#t|jdkr#|dd}|jd|}t|d}||fVWddS1sAwYdSd|vr|jrc|drZ|dd}|jd|}n%t|j}|jd d d |r||dd }|jd d ||}n||}t|d }||fVWddS1swYdStd|d)aA Open a file for reading or writing on local disk or URL to object store. Arg "key" should contain the extension for the file. Mode is either "r" or "w". Arg "add_subdir" can be set to True, in which case on the local filesystem it will create an extra subdirectory of ``self.storage_path`` with the first three letters of ``key``, preventing big datasets from exhausting the filesystem with one big directory. This arg is ignored for URLs. Yields a tuple of (open_file_object, path_or_url). rwz6Opening for both reading and writing is not supported.rarNrbT)exist_okparents)rowbz)Unsupported file mode (missing r or w): '') startswithlenr#rrer mkdir ValueError) rr$rirj input_pathf output_pathpsubdirrrr open_fileobjs6   "       "zFileIO.open_fileobj)F) r>r?r@rArrdrrCboolr tupler} __classcell__rrrgr r`s r`c TeZdZdZdZdeffdd Ze  dded e d e e d e j fd d Z ZS)LilcomFilesReadera Reads Lilcom-compressed files from a directory on the local filesystem, or a URL to supported object store (S3, AIStore, etc.). ``storage_path`` corresponds to the directory path; ``storage_key`` for each utterance is the name of the file in that directory. lilcom_filesr#ctt||_dSrrcrdr`iorr#r;r<rgrr rd7 zLilcomFilesReader.__init__rNr$rIrJrcCsN|jj|dd\}}t|}Wdn1swY|||S)Nrkri)rr}lilcom decompressrLrr$rIrJryrxarrrrr rL; zLilcomFilesReader.readrMr>r?r@rAr!rrdrrCrFr rDrErLrrrrgr r, rcsZeZdZdZdZddedeffdd Zede fd d Z d e d e j de fd dZ ZS)LilcomFilesWritera Writes Lilcom-compressed files to a directory on the local filesystem, or a URL to supported object store (S3, AIStore, etc.). ``storage_path`` corresponds to the directory path; ``storage_key`` for each utterance is the name of the file in that directory. rr# tick_powercstt||_||_dSr)rcrdr`rr)rr#rr;r<rgrr rdRs   zLilcomFilesWriter.__init__rcC|jjSrrr#rrrr r#WzLilcomFilesWriter.storage_pathr$r%cCs|ds |d}tj||jd}|jj|ddd(\}}|||jjs:dt |j dd}Wd|SWd|S1sEwY|S)Nz.llcrrlTrjra) rfrcompressrrr}r'rejoinr parts)rr$r%serialized_featsryrzrrr r'[s    zLilcomFilesWriter.write)r)r>r?r@rAr!rrFrdrBrCr#rDrEr'rrrrgr rGs rc r)NumpyFilesReadera$ Reads non-compressed numpy arrays from files in a directory on the local filesystem, or a URL to supported object store (S3, AIStore, etc.). ``storage_path`` corresponds to the directory path; ``storage_key`` for each utterance is the name of the file in that directory. numpy_filesr#crrrrrgrr rdvrzNumpyFilesReader.__init__rNr$rIrJrcCsN|jj|dd\}}tj|dd}Wdn1swY|||S)NrkrF allow_pickle)rr}rDloadrrrr rLzrzNumpyFilesReader.readrMrrrrgr rkrrcsTeZdZdZdZdeffdd ZedefddZ d ed e j defd d Z Z S) NumpyFilesWritera# Writes non-compressed numpy arrays to files in a directory on the local filesystem, or a URL to supported object store (S3, AIStore, etc.). ``storage_path`` corresponds to the directory path; ``storage_key`` for each utterance is the name of the file in that directory. rr#crrrrrgrr rdrzNumpyFilesWriter.__init__rcCrrrrrrr r#rzNumpyFilesWriter.storage_pathr$r%cCs|ds |d}|jj|ddd+\}}tj||dd|jjs5dt|jdd}Wd|SWd|S1s@wY|S) Nz.npyrlTrFrrar) rfrr}rDsavererr r)rr$r%ryrzrrr r's   zNumpyFilesWriter.write)r>r?r@rAr!rrdrBrCr#rDrEr'rrrrgr rs rcCtdstddS)Nh5pyzETo read and write HDF5 file formats, please 'pip install h5py' first.rrwrrrr check_h5py_installed r)maxsizer#cCstddl}||dS)aq Helper internal function used in HDF5 readers. It opens the HDF files and keeps their handles open in a global program cache to avoid excessive amount of syscalls when the Reader class is instantiated and destroyed in a loop repeatedly (frequent use-case). The file handles can be freed at any time by calling ``close_cached_file_handles()``. rNrk)rrFile)r#rrrr lookup_cache_or_opens  rcCs |tdS)z Helper internal function to retrieve the chunk size from an HDF5 file. Helps avoid unnecessary repeated disk reads. r)CHUNK_SIZE_KEY)h5_file_handlerrr lookup_chunk_sizes rcCstttdS)z Closes the cached file handles in ``lookup_cache_or_open`` and ``lookup_reader_cache_or_open`` (see respective docs for more details). N)r cache_clearrlookup_reader_cache_or_openrrrr close_cached_file_handless rc r)NumpyHdf5Readera[ Reads non-compressed numpy arrays from a HDF5 file with a "flat" layout. Each array is stored as a separate HDF ``Dataset`` because their shapes (numbers of frames) may vary. ``storage_path`` corresponds to the HDF5 file path; ``storage_key`` for each utterance is the key corresponding to the array (i.e. HDF5 "Group" name). numpy_hdf5r#crrrcrdrhdfrrgrr rdrzNumpyHdf5Reader.__init__rNr$rIrJrcCs|j|||Sr)rrKrrr rLs zNumpyHdf5Reader.readrMrrrrgr rrrcsteZdZdZdZddedeffdd Zedefd d Z d ed e j defd dZ dddZ ddZddZZS)NumpyHdf5Writera  Writes non-compressed numpy arrays to a HDF5 file with a "flat" layout. Each array is stored as a separate HDF ``Dataset`` because their shapes (numbers of frames) may vary. ``storage_path`` corresponds to the HDF5 file path; ``storage_key`` for each utterance is the key corresponding to the array (i.e. HDF5 "Group" name). Internally, this class opens the file lazily so that this object can be passed between processes without issues. This simplifies the parallel feature extraction code. rrlr#ricsVttddl}t|}||jdkr|jdnd|_|j|j |d|_ dS)a :param storage_path: Path under which we'll create the HDF5 file. We will add a ``.h5`` suffix if it is not already in ``storage_path``. :param mode: Modes supported by h5py: w Create file, truncate if exists (default) w- or x Create file, fail if exists a Read/write if exists, create otherwise rN.h5r) rcrdrrr with_suffixsuffix storage_path_rr#r)rr#rir;r<rr{rgrr rds zNumpyHdf5Writer.__init__rcC t|jSrrCrrrrr r# zNumpyHdf5Writer.storage_pathr$r%cCs|jj||d|S)Ndata)rcreate_datasetr&rrr r'szNumpyHdf5Writer.writeNcC |jSrrcloserrrr r zNumpyHdf5Writer.closecCr7rrrrrr r8r9zNumpyHdf5Writer.__enter__cC |dSrrrexc_typeexc_valexc_tbrrr r= zNumpyHdf5Writer.__exit__)rlrN)r>r?r@rAr!rrCrdrBr#rDrEr'rr8r=rrrrgr rs  rc r)LilcomHdf5Readera^ Reads lilcom-compressed numpy arrays from a HDF5 file with a "flat" layout. Each array is stored as a separate HDF ``Dataset`` because their shapes (numbers of frames) may vary. ``storage_path`` corresponds to the HDF5 file path; ``storage_key`` for each utterance is the key corresponding to the array (i.e. HDF5 "Group" name). lilcom_hdf5r#crrrrrgrr rd0rzLilcomHdf5Reader.__init__rNr$rIrJrcCs$t|j|d}|||S)Nr)rrrtobytesrr$rIrJrrrr rL4s zLilcomHdf5Reader.readrMrrrrgr r%rrcs|eZdZdZdZ  ddededeffdd Ze d efd d Z d ede j d efddZ dddZddZddZZS)LilcomHdf5Writera] Writes lilcom-compressed numpy arrays to a HDF5 file with a "flat" layout. Each array is stored as a separate HDF ``Dataset`` because their shapes (numbers of frames) may vary. ``storage_path`` corresponds to the HDF5 file path; ``storage_key`` for each utterance is the key corresponding to the array (i.e. HDF5 "Group" name). rrrlr#rrics\ttddl}t|}||jdkr|jdnd|_|j|j |d|_ ||_ dS)a :param storage_path: Path under which we'll create the HDF5 file. We will add a ``.h5`` suffix if it is not already in ``storage_path``. :param tick_power: Determines the lilcom compression accuracy; the input will be compressed to integer multiples of 2^tick_power. :param mode: Modes supported by h5py: w Create file, truncate if exists (default) w- or x Create file, fail if exists a Read/write if exists, create otherwise rNrr) rcrdrrr rrrrr#rr)rr#rrir;r<rr{rgrr rdNs  zLilcomHdf5Writer.__init__rcCrrrrrrr r#krzLilcomHdf5Writer.storage_pathr$r%cCs*tj||jd}|jj|t|d|S)Nrr)rrrrrrDvoid)rr$r%rrrr r'oszLilcomHdf5Writer.writeNcCrrrrrrr rtrzLilcomHdf5Writer.closecCr7rrrrrr r8wr9zLilcomHdf5Writer.__enter__cCrrrrrrr r=zrzLilcomHdf5Writer.__exit__)rrlrr>r?r@rAr!rrFrCrdrBr#rDrEr'rr8r=rrrrgr rCs$ r__LHOTSE_INTERNAL_CHUNK_SIZE__c r)ChunkedLilcomHdf5Readera Reads lilcom-compressed numpy arrays from a HDF5 file with chunked lilcom storage. Each feature matrix is stored in an array of chunks - binary data compressed with lilcom. Upon reading, we check how many chunks need to be retrieved to avoid excessive I/O. ``storage_path`` corresponds to the HDF5 file path; ``storage_key`` for each utterance is the key corresponding to the array (i.e. HDF5 "Group" name). chunked_lilcom_hdf5r#crrrrrgrr rdrz ChunkedLilcomHdf5Reader.__init__rNr$rIrJrc Cst|j}t||}|durt||}nd}dd|j|||D}|r0tj|dd}ntg}||} || } |durF|| } nd} || | S)NcSsg|] }t|qSr)rrrr,rrrr s z0ChunkedLilcomHdf5Reader.read..raxis)rrrrrD concatenater2) rr$rIrJ chunk_sizeleft_chunk_idxright_chunk_idxdecompressed_chunksr shift_framesleft_offset_shiftright_offset_shiftrrr rLs"     zChunkedLilcomHdf5Reader.readrMrrrrgr rs  rc seZdZdZdZ   ddededed effd d Ze d efd dZ dede j d efddZ dddZddZddZZS)ChunkedLilcomHdf5Writera Writes lilcom-compressed numpy arrays to a HDF5 file with chunked lilcom storage. Each feature matrix is stored in an array of chunks - binary data compressed with lilcom. Upon reading, we check how many chunks need to be retrieved to avoid excessive I/O. ``storage_path`` corresponds to the HDF5 file path; ``storage_key`` for each utterance is the key corresponding to the array (i.e. HDF5 "Group" name). rrdrlr#rrric sttddl}t|}||jdkr|jdnd|_||_||_ |j |j |d|_ t |j vrM|j t d} | t ksKJd|j d| ddS|j jt |j d dS) a :param storage_path: Path under which we'll create the HDF5 file. We will add a ``.h5`` suffix if it is not already in ``storage_path``. :param tick_power: Determines the lilcom compression accuracy; the input will be compressed to integer multiples of 2^tick_power. :param chunk_size: How many frames to store per chunk. Too low a number will require many reads for long feature matrices, too high a number will require to read more redundant data. :param mode: Modes supported by h5py: w Create file, truncate if exists (default) w- or x Create file, fail if exists a Read/write if exists, create otherwise rNrrrz*Error: attempted to write with chunk size z2 to an h5py file that was created with chunk size .r)rcrdrrr rrrrrrr#rrr) rr#rrrir;r<rr{retrieved_chunk_sizergrr rds$    z ChunkedLilcomHdf5Writer.__init__rcCrrrrrrr r#rz$ChunkedLilcomHdf5Writer.storage_pathr$r%c Csztddl}ddlm}|||j|jd}|jj||t dt |fd}t |D]\}}t j |t jd||<q,|S)Nrlilcom_compress_chunkedrruint8)dtyper1)r)rrlhotse.features.compressionrrrrr vlen_dtyperDrru enumerate frombufferr) rr$r%rrrdsetidxfeatrrr r's  zChunkedLilcomHdf5Writer.writeNcCrrrrrrr r rzChunkedLilcomHdf5Writer.closecCr7rrrrrr r8 r9z!ChunkedLilcomHdf5Writer.__enter__cCrrrrrrr r=rz ChunkedLilcomHdf5Writer.__exit__)rrrlrrrrrgr rs* * ric sXeZdZdZdZeZdeffdd Ze  dde d e d e e d e jfd d ZZS)LilcomChunkyReadera` Reads lilcom-compressed numpy arrays from a binary file with chunked lilcom storage. Each feature matrix is stored in an array of chunks - binary data compressed with lilcom. Upon reading, we check how many chunks need to be retrieved to avoid excessive I/O. ``storage_path`` corresponds to the binary file path. ``storage_key`` for each utterance is a comma separated list of offsets in the file. The first number is the offset for the whole array, and the following numbers are relative offsets for each chunk. These offsets are relative to the previous chunk start. lilcom_chunkyr#cst||_dSr)rcrdr#rrgrr rd/s  zLilcomChunkyReader.__init__rNr$rIrJrcCst||j}|durt||jd}nd}ttt|d}t|}|||}g}t |j d }t |D]\} } | | | || | q9Wdn1sWwYdd|D} | rmtj| dd} ntg} |j|} || }|dur|| }nd}| ||S)Nrm,rncSsg|]}t|qSrrrrrrr rMsz+LilcomChunkyReader.read..rr)r CHUNK_SIZErr4maprFsplitrDcumsumopenr#rseekappendrLrr2)rr$rIrJrr chunk_offsets chunk_datafileoffsetendrrrrrrrr rL3s0       zLilcomChunkyReader.readrM)r>r?r@rAr!CHUNKY_FORMAT_CHUNK_SIZErrrdrrCrFr rDrErLrrrrgr rs" rcseZdZdZdZeZ  ddedede ffdd Z e d e fd d Z d e de jd e fddZdddZddZddZZS)LilcomChunkyWritera_ Writes lilcom-compressed numpy arrays to a binary file with chunked lilcom storage. Each feature matrix is stored in an array of chunks - binary data compressed with lilcom. Upon reading, we check how many chunks need to be retrieved to avoid excessive I/O. ``storage_path`` corresponds to the binary file path. ``storage_key`` for each utterance is a comma separated list of offsets in the file. The first number is the offset for the whole array, and the following numbers are relative offsets for each chunk. These offsets are relative to the previous chunk start. rrrrr#rricsttd|vr |d}|dvsJt|}||jdkr#|jdnd|_||_t|j|d|_ |j |_ dS)a) :param storage_path: Path under which we'll create the binary file. :param tick_power: Determines the lilcom compression accuracy; the input will be compressed to integer multiples of 2^tick_power. :param chunk_size: How many frames to store per chunk. Too low a number will require many reads for long feature matrices, too high a number will require to read more redundant data. :param mode: Modes, one of: "w" (write) or "a" (append); can be "wb" and "ab", "b" is implicit b)rrabz.lcarN) rcrdr rrrrrr#rtell curr_offset)rr#rrir;r<r{rgrr rdqs  zLilcomChunkyWriter.__init__rcCrrrrrrr r#rzLilcomChunkyWriter.storage_pathr$r%c Cslddlm}|||j|jd}|jg}t|D]\}}|j|}|||j|7_qd t t |S)Nrrrr) rrrrr rrr'rrrrC) rr$r%rroffsetsrrnbytesrrr r's    zLilcomChunkyWriter.writeNcCs|jdSr)rrrrrr rzLilcomChunkyWriter.closecCr7rrrrrr r8r9zLilcomChunkyWriter.__enter__cCrrrrrrr r=rzLilcomChunkyWriter.__exit__)rrrr)r>r?r@rAr!rrrrFrCrdrBr#rDrEr'rr8r=rrrrgr r_s&   rc sNeZdZdZdZfddZe  d deded e ed e j fd d Z Z S)LilcomURLReaderao Downloads Lilcom-compressed files from a URL (S3, GCP, Azure, HTTP, etc.). ``storage_path`` corresponds to the root URL (e.g. "s3://my-data-bucket") ``storage_key`` will be concatenated to ``storage_path`` to form a full URL (e.g. "my-feature-file.llc") .. caution:: Requires ``smart_open`` to be installed (``pip install smart_open``). lilcom_urlctt|i||_dSr)rcrdr_innerr:rgrr rd zLilcomURLReader.__init__rNr$rIrJrcCs|j|||Sr)rrLrKrrr rLszLilcomURLReader.readrM)r>r?r@rAr!rdrrCrFr rDrErLrrrrgr r s  r csNeZdZdZdZfddZedefddZded e j defd d Z Z S) LilcomURLWriteraj Writes Lilcom-compressed files to a URL (S3, GCP, Azure, HTTP, etc.). ``storage_path`` corresponds to the root URL (e.g. "s3://my-data-bucket") ``storage_key`` will be concatenated to ``storage_path`` to form a full URL (e.g. "my-feature-file.llc") .. caution:: Requires ``smart_open`` to be installed (``pip install smart_open``). rcrr)rcrdrrr:rgrr rdrzLilcomURLWriter.__init__rcCrr)rr#rrrr r#rzLilcomURLWriter.storage_pathr$r%cCs|j||Sr)rr'r&rrr r'r zLilcomURLWriter.write) r>r?r@rAr!rdrBrCr#rDrEr'rrrrgr rs  rcCr)Nkaldi_native_ioDTo read Kaldi feats.scp, please 'pip install kaldi_native_io' first.rrrrr check_kaldi_native_io_installedrrcCstddl}|d|S)ar Helper internal function used in KaldiReader. It opens kaldi scp files and keeps their handles open in a global program cache to avoid excessive amount of syscalls when the Reader class is instantiated and destroyed in a loop repeatedly (frequent use-case). The file handles can be freed at any time by calling ``close_cached_file_handles()``. rNzscp:)rrRandomAccessFloatMatrixReader)r#rrrr rs rc r) KaldiReadera& Reads Kaldi's "feats.scp" file using kaldi_native_io. ``storage_path`` corresponds to the path to ``feats.scp``. ``storage_key`` corresponds to the utterance-id in Kaldi. .. caution:: Requires ``kaldi_native_io`` to be installed (``pip install kaldi_native_io``). kaldiior#csJt||_|drt|j|_dStddl}d|_|j|_ dS)Nz.scpr) rcrdr#rfrstoragerr FloatMatrixreader)rr#r;r<rrgrr rds   zKaldiReader.__init__rNr$rIrJrcCs:|jdurt|j|}n |j|j}|||Sr)rrDcopyrrLr#numpyrrrr rLs  zKaldiReader.readrMrrrrgr rs   rcsveZdZdZdZ ddedeffdd Zede fd d Z d e d e j de fd dZ dddZddZddZZS) KaldiWritera Write data to Kaldi's "feats.scp" and "feats.ark" files using kaldi_native_io. ``storage_path`` corresponds to a directory where we'll create "feats.scp" and "feats.ark" files. ``storage_key`` corresponds to the utterance-id in Kaldi. The following ``compression_method`` values are supported by kaldi_native_io:: kAutomaticMethod = 1 kSpeechFeature = 2 kTwoByteAuto = 3 kTwoByteSignedInteger = 4 kOneByteAuto = 5 kOneByteUnsignedInteger = 6 kOneByteZeroOne = 7 .. note:: Setting compression_method works only with 2D arrays. Example:: >>> data = np.random.randn(131, 80) >>> with KaldiWriter('featdir') as w: ... w.write('utt1', data) >>> reader = KaldiReader('featdir/feats.scp') >>> read_data = reader.read('utt1') >>> np.testing.assert_equal(data, read_data) .. caution:: Requires ``kaldi_native_io`` to be installed (``pip install kaldi_native_io``). rrmr#compression_methodcsztdstdddl}tt||_|jjdddt|jd|_ | d|jd|jd |_ | ||_ dS) NrrrT)rproz feats.scpzark,scp:z /feats.ark,z /feats.scp)rrwrrcrdr storage_dirrvrCrCompressedMatrixWriterrCompressionMethodr)rr#rr;r<rrgrr rdQs  zKaldiWriter.__init__rcCs|jSr)rrrrr r#gszKaldiWriter.storage_pathr$r%cCs|j|||j|Sr)rr'rr&rrr r'kszKaldiWriter.writeNcCrr)rrrrrr rorzKaldiWriter.closecCr7rrrrrr r8rr9zKaldiWriter.__enter__cCrrrrrrr r=urzKaldiWriter.__exit__)rmr)r>r?r@rAr!rrFrdrBrCr#rDrEr'rr8r=rrrrgr r.s rr/cCsd|vSNmemoryr)r/rrr is_in_memory~sr%cCsd|vsJt|Sr#)r_r\rrr get_memory_writers r&c @FeZdZdZdZddZe  d deded e ed e j fd d Z dS)MemoryLilcomReader memory_lilcomcOrrrr:rrr rdr9zMemoryLilcomReader.__init__rNraw_datarIrJrcCt|}|||Srrrr+rIrJrrrr rL  zMemoryLilcomReader.readrM r>r?r@rAr!rdrbytesrFr rDrErLrrrr r( r(c@sleZdZdZdZdddeddfdd Zedd d Zd e d e j de fddZ dddZddZddZdS)MemoryLilcomWriterr)r*rlilcom_tick_powerr4rNcOs ||_dSrr3)rr4r;r<rrr rds zMemoryLilcomWriter.__init__cCrrrrrrr r#r(zMemoryLilcomWriter.storage_pathr$r%cCs(t|jtjs Jdtj||jdS)Nz7Lilcom compression supports only floating-point arrays.r)rD issubdtyperfloatingrrr4r&rrr r's zMemoryLilcomWriter.writecCrrrrrrr rr9zMemoryLilcomWriter.closecCr7rrrrrr r8r9zMemoryLilcomWriter.__enter__cCrrrrrrr r=r9zMemoryLilcomWriter.__exit__r)r>r?r@rAr!rFrdrBr#rCrDrEr0r'rr8r=rrrr r2s    r2c @r')MemoryRawReaderr) memory_rawcOrrrr:rrr rdr9zMemoryRawReader.__init__rNr+rIrJrcCr,r)pickleloadsr-rrr rLr.zMemoryRawReader.readrMr/rrrr r7r1r7c@s\eZdZdZdZddZedddZd ed e j de fd d Z dd dZ ddZddZdS)MemoryRawWriterr)r8cOrrrr:rrr rdr9zMemoryRawWriter.__init__rNcCrrrrrrr r#r(zMemoryRawWriter.storage_pathr$r%cCs t|Sr)r9dumpsr&rrr r'rzMemoryRawWriter.writecCrrrrrrr rr9zMemoryRawWriter.closecCr7rrrrrr r8r9zMemoryRawWriter.__enter__cCrrrrrrr r=r9zMemoryRawWriter.__exit__r)r>r?r@rAr!rdrBr#rCrDrEr0r'rr8r=rrrr r;s   r;c @r')MemoryNpyReaderr) memory_npycOrrrr:rrr rdr9zMemoryNpyReader.__init__rNr+rIrJrcCst|}t|}|||Sr)rrDr)rr+rIrJstreamrrrr rLs  zMemoryNpyReader.readrMr/rrrr r=r1r=c@s,eZdZdZdZddZdejfddZdS) DummySharReaderr)sharcOrrrr:rrr rdr9zDummySharReader.__init__rcOstd)NzeInconsistent state: found a Lhotse Shar placeholder array that was not filled during deserialization.) RuntimeErrorr:rrr rL szDummySharReader.readN) r>r?r@rAr!rdrDrErLrrrr r@s r@r)Sr9abcrr contextlibr functoolsrrrmathrrpathlibr typingr r r r rrrrD lhotse.arrayrrlhotse.cachingrlhotse.serializationr lhotse.utilsrrrrrrrrHrPrRrCrSrWrYr]r_r`rrrrrrrFrrrrrrrrrrrrr rrrrrr~r%r&r(r2r7r;r=r@rrrr s        m(  <  2:5VAQ'J!