codingutf8 2023 the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache orglicenseslicense2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license feature extractor class for tvlt from math import ceil from typing import list optional union import numpy as np from audioutils import melfilterbank spectrogram windowfunction from featureextractionsequenceutils import batchfeature sequencefeatureextractor from utils import tensortype logging logger logging getloggername class tvltfeatureextractorsequencefeatureextractor r constructs a tvlt audio feature extractor this feature extractor can be used to prepare audios for the model this feature extractor inherits from featureextractionmixin which contains most of the main methods users should refer to this superclass for more information regarding those methods args spectrogramlength dictstr int optional defaults to 2048 the time length of each audio spectrogram numchannels int optional defaults to 1 number of audio channels patchsize listint optional defaults to 16 16 the patch size of audio patch embedding featuresize int optional defaults to 128 the frequency length of audio spectrogram samplingrate int optional defaults to 44100 the sampling rate at which the audio files should be digitalized expressed in hertz hz hoplengthtosamplingrate int optional defaults to 86 hop length is length of the overlaping windows for the stft used to obtain the mel frequency coefficients for example with sampling rate 44100 the hop length is 512 with 44100 512 86 nfft int optional defaults to 2048 size of the fourier transform paddingvalue float optional defaults to 0 0 padding value used to pad the audio should correspond to silences compute the logmel spectrogram of the provided audio gives similar results to whisper s original torch implementation with 1e5 tolerance main method to prepare one or several audios for the model args rawspeech np ndarray listfloat listnp ndarray listlistfloat the sequence or batch of sequences to be padded each sequence can be a numpy array a list of float values a list of numpy arrays or a list of list of float values must be mono channel audio not stereo i e single float per timestep returntensors str or utils tensortype optional if set will return tensors instead of list of python integers acceptable values are pt return pytorch torch tensor objects np return numpy np ndarray objects returnattentionmask bool optional default to true whether to return the attention mask if left to the default will return the attention mask according to the specific featureextractor s default what are attention masks glossaryattentionmask tip for tvlttransformer models attentionmask should alwys be passed for batched inference to avoid subtle bugs tip samplingrate int optional the sampling rate at which the rawspeech input was sampled it is strongly recommended to pass samplingrate at the forward call to prevent silent errors and allow automatic speech recognition pipeline current model supports sampling rate 16000 and 44100 resample bool optional defaults to false if the sampling rate is not matched resample the input audio to match maskaudio bool optional defaults to false whether or not to mask input audio for mae task returns batchfeature a batchfeature with the following fields audiovalues audio values to be fed to a model of shape batchsize numchannels height width audiomask audio masks to be fed to a model of shape batchsize numaudiopatches always return batch convert audio signals to log mel spectrograms truncate by time axis create audio attention mask convert into correct format for padding return as batchfeature coding utf 8 2023 the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache org licenses license 2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license feature extractor class for tvlt constructs a tvlt audio feature extractor this feature extractor can be used to prepare audios for the model this feature extractor inherits from featureextractionmixin which contains most of the main methods users should refer to this superclass for more information regarding those methods args spectrogram_length dict str int optional defaults to 2048 the time length of each audio spectrogram num_channels int optional defaults to 1 number of audio channels patch_size list int optional defaults to 16 16 the patch size of audio patch embedding feature_size int optional defaults to 128 the frequency length of audio spectrogram sampling_rate int optional defaults to 44100 the sampling rate at which the audio files should be digitalized expressed in hertz hz hop_length_to_sampling_rate int optional defaults to 86 hop length is length of the overlaping windows for the stft used to obtain the mel frequency coefficients for example with sampling rate 44100 the hop length is 512 with 44100 512 86 n_fft int optional defaults to 2048 size of the fourier transform padding_value float optional defaults to 0 0 padding value used to pad the audio should correspond to silences compute the log mel spectrogram of the provided audio gives similar results to whisper s original torch implementation with 1e 5 tolerance main method to prepare one or several audio s for the model args raw_speech np ndarray list float list np ndarray list list float the sequence or batch of sequences to be padded each sequence can be a numpy array a list of float values a list of numpy arrays or a list of list of float values must be mono channel audio not stereo i e single float per timestep return_tensors str or utils tensortype optional if set will return tensors instead of list of python integers acceptable values are pt return pytorch torch tensor objects np return numpy np ndarray objects return_attention_mask bool optional default to true whether to return the attention mask if left to the default will return the attention mask according to the specific feature_extractor s default what are attention masks glossary attention mask tip for tvlttransformer models attention_mask should alwys be passed for batched inference to avoid subtle bugs tip sampling_rate int optional the sampling rate at which the raw_speech input was sampled it is strongly recommended to pass sampling_rate at the forward call to prevent silent errors and allow automatic speech recognition pipeline current model supports sampling rate 16000 and 44100 resample bool optional defaults to false if the sampling rate is not matched resample the input audio to match mask_audio bool optional defaults to false whether or not to mask input audio for mae task returns batchfeature a batchfeature with the following fields audio_values audio values to be fed to a model of shape batch_size num_channels height width audio_mask audio masks to be fed to a model of shape batch_size num_audio_patches always return batch convert audio signals to log mel spectrograms truncate by time axis create audio attention mask the maximum number of audio patches in a batch convert into correct format for padding the maximum audio size in a batch return as batchfeature
from math import ceil
from typing import List, Optional, Union

import numpy as np

from ...audio_utils import mel_filter_bank, spectrogram, window_function
from ...feature_extraction_sequence_utils import BatchFeature, SequenceFeatureExtractor
from ...utils import TensorType, logging


logger = logging.get_logger(__name__)


class TvltFeatureExtractor(SequenceFeatureExtractor):
    r

    model_input_names = ["audio_values", "audio_mask"]

    def __init__(
        self,
        spectrogram_length=2048,
        num_channels=1,
        patch_size=[16, 16],
        feature_size=128,
        sampling_rate=44100,
        hop_length_to_sampling_rate=86,
        n_fft=2048,
        padding_value=0.0,
        **kwargs,
    ):
        super().__init__(
            feature_size=feature_size,
            sampling_rate=sampling_rate,
            padding_value=padding_value,
            **kwargs,
        )

        self.spectrogram_length = spectrogram_length
        self.num_channels = num_channels
        self.patch_size = patch_size
        self.freq_len = feature_size // self.patch_size[1]
        self.n_fft = n_fft
        self.hop_length = sampling_rate // hop_length_to_sampling_rate
        self.sampling_rate = sampling_rate
        self.padding_value = padding_value
        self.mel_filters = mel_filter_bank(
            num_frequency_bins=1 + n_fft // 2,
            num_mel_filters=feature_size,
            min_frequency=0.0,
            max_frequency=22050.0,
            sampling_rate=sampling_rate,
            norm="slaney",
            mel_scale="slaney",
        ).T

    def _np_extract_fbank_features(self, waveform: np.array) -> np.ndarray:
        
        log_spec = spectrogram(
            waveform,
            window_function(self.n_fft, "hann"),
            frame_length=self.n_fft,
            hop_length=self.hop_length,
            power=2.0,
            mel_filters=self.mel_filters.T,
            log_mel="dB",
            db_range=80.0,
        )
        log_spec = log_spec[:, :-1]
        log_spec = log_spec - 20.0
        log_spec = np.clip(log_spec / 40.0, -2.0, 0.0) + 1.0
        return log_spec

    def __call__(
        self,
        raw_speech: Union[np.ndarray, List[float], List[np.ndarray], List[List[float]]],
        return_tensors: Optional[Union[str, TensorType]] = None,
        return_attention_mask: Optional[bool] = True,
        sampling_rate: Optional[int] = None,
        resample: bool = False,
        mask_audio: bool = False,
        **kwargs,
    ) -> BatchFeature:
        

        if sampling_rate is not None:
            if sampling_rate != self.sampling_rate:
                raise ValueError(
                    "This feature extractor is set to support sampling rate"
                    f" of {self.sampling_rate}. Please make sure that the provided `raw_speech` input was sampled"
                    f" with {self.sampling_rate} and not {sampling_rate}."
                )
        else:
            logger.warning(
                "It is strongly recommended to pass the `sampling_rate` argument to this function. "
                "Failing to do so can result in silent errors that might be hard to debug."
            )

        is_batched_numpy = isinstance(raw_speech, np.ndarray) and len(raw_speech.shape) > 1
        if is_batched_numpy and len(raw_speech.shape) > 2:
            raise ValueError(f"Only mono-channel audio is supported for input to {self}")
        is_batched = is_batched_numpy or (
            isinstance(raw_speech, (list, tuple)) and (isinstance(raw_speech[0], (np.ndarray, tuple, list)))
        )
        if is_batched:
            raw_speech = [np.asarray([speech], dtype=np.float32).T for speech in raw_speech]
        elif not is_batched and not isinstance(raw_speech, np.ndarray):
            raw_speech = np.asarray(raw_speech, dtype=np.float32)
        elif isinstance(raw_speech, np.ndarray) and raw_speech.dtype is np.dtype(np.float64):
            raw_speech = raw_speech.astype(np.float32)
        
        if not is_batched:
            raw_speech = [np.asarray([raw_speech]).T]

        
        audio_features = [
            self._np_extract_fbank_features(waveform.squeeze()).T[: self.spectrogram_length] for waveform in raw_speech
        ]
        if isinstance(audio_features[0], List):
            audio_features = [np.asarray(feature, dtype=np.float32) for feature in audio_features]

        
        max_patch_len = max(
            [ceil(feature.shape[0] / self.patch_size[0]) * self.freq_len for feature in audio_features]
        )  
        if return_attention_mask:
            audio_mask = [
                (ceil(feature.shape[0] / self.patch_size[0]) * self.freq_len) * [1]
                + (max_patch_len - ceil(feature.shape[0] / self.patch_size[0]) * self.freq_len) * [0]
                for feature in audio_features
            ]
            audio_mask = np.array(audio_mask).astype(np.float32)

        
        max_time_len = max_patch_len // self.freq_len * self.patch_size[0]  
        padded_audio_features = np.ones([len(audio_features), 1, max_time_len, self.feature_size]).astype(np.float32)
        padded_audio_features = padded_audio_features * self.padding_value
        for i in range(len(audio_features)):
            feature = audio_features[i]
            padded_audio_features[i, :, : feature.shape[0], :] = feature

        
        if return_attention_mask:
            data = {"audio_values": padded_audio_features, "audio_mask": audio_mask}
        else:
            data = {"audio_values": padded_audio_features}

        encoded_inputs = BatchFeature(data=data, tensor_type=return_tensors)
        return encoded_inputs