a h! @sddlZddlmZddlmZmZddlZddlmZddl m Z ddl m Z ddl mZdd d d d gd dddddd d gd ddddddd gd dddZGddde ZdS)N)Path)TupleUnion)Tensor)Dataset)download_url_to_file) _extract_tarZTEDLIUM_release1z:http://www.openslr.org/resources/7/TEDLIUM_release1.tar.gzZ@30301975fd8c5cac4040c261c0852f57cfa8adbbad2ce78e77e4986957445f27train)r testdevzTEDLIUM.150K.dic)folder_in_archiveurlchecksum data_pathsubsetsupported_subsetsdictZTEDLIUM_release2z;http://www.openslr.org/resources/19/TEDLIUM_release2.tar.gzZ@93281b5fcaaae5c88671c9d000b443cb3c7ea3499ad12010b3934ca41a7b9c58zTEDLIUM.152k.diczTEDLIUM_release-3z9http://www.openslr.org/resources/51/TEDLIUM_release-3.tgzZ@ad1e454d14d1ad550bc2564c462d87c7a7ec83d4dc2b9210f22ab4973b9eccdbzdata/)release1Zrelease2release3c @seZdZdZdeeefeeeedddd Zee ee e e ee e e fd d d Z dee e e e e gdddZe e e e ee e e fdddZe dddZeddZdS)TEDLIUMa*Tedlium* :cite:`rousseau2012tedlium` dataset (releases 1,2 and 3). Args: root (str or Path): Path to the directory where the dataset is found or downloaded. release (str, optional): Release version. Allowed values are ``"release1"``, ``"release2"`` or ``"release3"``. (default: ``"release1"``). subset (str, optional): The subset of dataset to use. Valid options are ``"train"``, ``"dev"``, and ``"test"``. Defaults to ``"train"``. download (bool, optional): Whether to download the dataset if it is not found at root path. (default: ``False``). audio_ext (str, optional): extension for audio file (default: ``".sph"``) rr F.sphN)rootreleaserdownload audio_extreturnc sD||_|tvr@t|d}t|d}|r2|n t|d}ntd|t|t|dvr|td|t|dt|}tj|}tj ||} | dd}|d kr|d krtj ||t|d |_ ntj ||d ||_ ntj ||t|d ||_ |rTtj |j svtj | sJt|d } t|| | dt| n"tj|j svtd|j dg|_tj |j d} tt| D]drtj |j d} t| F} t| } dd|jfddt| DWdn1s0Yqtj ||t|d|_d|_dS)Nr rrzFThe release {} does not match any of the supported tedlium releases{} rzDThe subset {} does not match any of the supported tedlium subsets{} .rrr rlegacyr)Z hash_prefixz The path zT doesn't exist. Please check the ``root`` path or set `download=True` to download itstm.stmr c3s|]}|fVqdS)N).0linefiler!I/var/www/auris/lib/python3.9/site-packages/torchaudio/datasets/tedlium.py ~z#TEDLIUM.__init__..r) _ext_audio_RELEASE_CONFIGSkeys RuntimeErrorformatosfspathpathbasenamejoinsplit_pathisdirisfilerrexists _filelistsortedlistdirendswithopenlen readlinesreplaceextendrange _dict_path _phoneme_dict)selfrrrrrr rr1archiverZstm_pathflr!r$r&__init__:s`             BzTEDLIUM.__init__)fileidr#r0rcCstj|d|}t|d6}||}|dd\}}} } } } }Wdn1sX0Ytj|d|} |j| |j| | d\}}||||| | fS)aLoads a TEDLIUM dataset sample given a file name and corresponding sentence name. Args: fileid (str): File id to identify both text and audio files corresponding to the sample line (int): Line identifier for the sample inside the text file path (str): Dataset root path Returns: (Tensor, int, str, int, int, int): ``(waveform, sample_rate, transcript, talk_id, speaker_id, identifier)`` rr  NZsph) start_timeend_time)r.r0r2r<r>r3 _load_audior))rDrIr#r0Ztranscript_pathrFZ transcriptZtalk_id_Z speaker_idrLrM identifierZ wave_pathZwaveform sample_rater!r!r&_load_tedlium_items  8zTEDLIUM._load_tedlium_item>)r0rLrMrQrcCs@tt||}tt||}|||d}tj|fi|S)a"Default load function used in TEDLIUM dataset, you can overwrite this function to customize functionality and load individual sentences from a full ted audio talk file. Args: path (str): Path to audio file start_time (int): Time in seconds where the sample sentence stars end_time (int): Time in seconds where the sample sentence finishes sample_rate (float, optional): Sampling rate Returns: [Tensor, int]: Audio tensor representation and sample rate )Z frame_offsetZ num_frames)intfloat torchaudioload)rDr0rLrMrQkwargsr!r!r&rNs zTEDLIUM._load_audio)nrcCs|j|\}}||||jS)aLoad the n-th sample from the dataset. Args: n (int): The index of the sample to be loaded Returns: Tuple of the following items; Tensor: Waveform int: Sample rate str: Transcript int: Talk ID int: Speaker ID int: Identifier )r8rRr4)rDrYrIr#r!r!r& __getitem__szTEDLIUM.__getitem__)rcCs t|jS)zTEDLIUM dataset custom function overwritting len default behaviour. Returns: int: TEDLIUM dataset length )r=r8)rDr!r!r&__len__szTEDLIUM.__len__cCsz|jspi|_t|jdddD}|D]*}|}t|dd|j|d<q&Wdn1sf0Y|jS)zdict[str, tuple[str]]: Phonemes. Mapping from word to tuple of phonemes. Note that some words have empty phonemes. rzutf-8)encodingNr)rCr<rBr>stripr3tuplecopy)rDrFr#contentr!r!r& phoneme_dicts  :zTEDLIUM.phoneme_dict)rr Fr)rS)__name__ __module__ __qualname____doc__rstrrboolrHrTrrrRrUrNrZr[propertyrcr!r!r!r&r+s&  I$ r)r.pathlibrtypingrrrVZtorchrZtorch.utils.datarZtorchaudio._internalrZtorchaudio.datasets.utilsrr*rr!r!r!r&sB