Aluode/PerceptionLabPortable
0
1# coding=utf-82# Copyright 2022 The HuggingFace Inc. team.3#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8# http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15"""16Feature extractor class for Whisper17"""18 19from typing import Optional, Union20 21import numpy as np22 23from ... import is_torch_available24from ...audio_utils import mel_filter_bank, spectrogram, window_function25from ...feature_extraction_sequence_utils import SequenceFeatureExtractor26from ...feature_extraction_utils import BatchFeature27from ...utils import TensorType, logging28 29 30if is_torch_available():31 import torch32 33logger = logging.get_logger(__name__)34 35 36class WhisperFeatureExtractor(SequenceFeatureExtractor):37 r"""38 Constructs a Whisper feature extractor.39 40 This feature extractor inherits from [`~feature_extraction_sequence_utils.SequenceFeatureExtractor`] which contains41 most of the main methods. Users should refer to this superclass for more information regarding those methods.42 43 This class extracts mel-filter bank features from raw speech using a custom numpy implementation of the `Short Time44 Fourier Transform` which should match pytorch's `torch.stft` equivalent.45 46 Args:47 feature_size (`int`, *optional*, defaults to 80):48 The feature dimension of the extracted features.49 sampling_rate (`int`, *optional*, defaults to 16000):50 The sampling rate at which the audio files should be digitalized expressed in hertz (Hz).51 hop_length (`int`, *optional*, defaults to 160):52 Length of the overlapping windows for the STFT used to obtain the Mel Frequency coefficients.53 chunk_length (`int`, *optional*, defaults to 30):54 The maximum number of chunks of `sampling_rate` samples used to trim and pad longer or shorter audio55 sequences.56 n_fft (`int`, *optional*, defaults to 400):57 Size of the Fourier transform.58 padding_value (`float`, *optional*, defaults to 0.0):59 Padding value used to pad the audio. Should correspond to silences.60 dither (`float`, *optional*, defaults to 0.0):61 Adds dithering. In other words, adds a small Gaussian noise to each frame.62 E.g. use 0.0001 to add dithering with a normal distribution centered63 around 0.0 with standard deviation 0.0001 (assuming [-1,+1] range of raw_speech).64 The value 0.0 means no dithering.65 Dithering has similar effect as `spectrogram(mel_floor=...)`. It reduces66 the high log_mel_fbank values for signals with hard-zero sections,67 when VAD cutoff is present in the signal.68 """69 70 model_input_names = ["input_features"]71 72 def __init__(73 self,74 feature_size=80,75 sampling_rate=16000,76 hop_length=160,77 chunk_length=30,78 n_fft=400,79 padding_value=0.0,80 dither=0.0,81 return_attention_mask=False, # pad inputs to max length with silence token (zero) and no attention mask82 **kwargs,83 ):84 super().__init__(85 feature_size=feature_size,86 sampling_rate=sampling_rate,87 padding_value=padding_value,88 return_attention_mask=return_attention_mask,89 **kwargs,90 )91 self.n_fft = n_fft92 self.hop_length = hop_length93 self.chunk_length = chunk_length94 self.n_samples = chunk_length * sampling_rate95 self.nb_max_frames = self.n_samples // hop_length96 self.sampling_rate = sampling_rate97 self.dither = dither98 self.mel_filters = mel_filter_bank(99 num_frequency_bins=1 + n_fft // 2,100 num_mel_filters=feature_size,101 min_frequency=0.0,102 max_frequency=8000.0,103 sampling_rate=sampling_rate,104 norm="slaney",105 mel_scale="slaney",106 )107 108 def _np_extract_fbank_features(self, waveform_batch: np.ndarray, device: str) -> np.ndarray:109 """110 Compute the log-mel spectrogram of the provided audio, gives similar results to Whisper's original torch111 implementation with 1e-5 tolerance.112 """113 if device != "cpu":114 raise ValueError(115 f"Got device `{device}` for feature extraction, but feature extraction on CUDA accelerator "116 "devices requires torch, which is not installed. Either set `device='cpu'`, or "117 "install torch according to the official instructions: https://pytorch.org/get-started/locally/"118 )119 log_spec_batch = []120 for waveform in waveform_batch:121 log_spec = spectrogram(122 waveform,123 window_function(self.n_fft, "hann"),124 frame_length=self.n_fft,125 hop_length=self.hop_length,126 power=2.0,127 dither=self.dither,128 mel_filters=self.mel_filters,129 log_mel="log10",130 )131 log_spec = log_spec[:, :-1]132 log_spec = np.maximum(log_spec, log_spec.max() - 8.0)133 log_spec = (log_spec + 4.0) / 4.0134 log_spec_batch.append(log_spec)135 log_spec_batch = np.array(log_spec_batch)136 return log_spec_batch137 138 def _torch_extract_fbank_features(self, waveform: np.ndarray, device: str = "cpu") -> np.ndarray:139 """140 Compute the log-mel spectrogram of the audio using PyTorch's GPU-accelerated STFT implementation with batching,141 yielding results similar to cpu computing with 1e-5 tolerance.142 """143 waveform = torch.from_numpy(waveform).to(device, torch.float32)144 window = torch.hann_window(self.n_fft, device=device)145 146 # Note: it would be better to dither the chunked waveform,147 # so overlapping signal does not get the same dithering.148 # But, chunking is happening inside pytorch, so it is here.149 if self.dither != 0.0:150 waveform += self.dither * torch.randn(waveform.shape, dtype=waveform.dtype, device=waveform.device)151 152 stft = torch.stft(waveform, self.n_fft, self.hop_length, window=window, return_complex=True)153 magnitudes = stft[..., :-1].abs() ** 2154 155 mel_filters = torch.from_numpy(self.mel_filters).to(device, torch.float32)156 mel_spec = mel_filters.T @ magnitudes157 158 log_spec = torch.clamp(mel_spec, min=1e-10).log10()159 if waveform.dim() == 2:160 max_val = log_spec.max(dim=2, keepdim=True)[0].max(dim=1, keepdim=True)[0]161 log_spec = torch.maximum(log_spec, max_val - 8.0)162 else:163 log_spec = torch.maximum(log_spec, log_spec.max() - 8.0)164 log_spec = (log_spec + 4.0) / 4.0165 if device != "cpu":166 log_spec = log_spec.detach().cpu()167 return log_spec.numpy()168 169 @staticmethod170 # Copied from transformers.models.wav2vec2.feature_extraction_wav2vec2.Wav2Vec2FeatureExtractor.zero_mean_unit_var_norm171 def zero_mean_unit_var_norm(172 input_values: list[np.ndarray], attention_mask: list[np.ndarray], padding_value: float = 0.0173 ) -> list[np.ndarray]:174 """175 Every array in the list is normalized to have zero mean and unit variance176 """177 if attention_mask is not None:178 attention_mask = np.array(attention_mask, np.int32)179 normed_input_values = []180 181 for vector, length in zip(input_values, attention_mask.sum(-1)):182 normed_slice = (vector - vector[:length].mean()) / np.sqrt(vector[:length].var() + 1e-7)183 if length < normed_slice.shape[0]:184 normed_slice[length:] = padding_value185 186 normed_input_values.append(normed_slice)187 else:188 normed_input_values = [(x - x.mean()) / np.sqrt(x.var() + 1e-7) for x in input_values]189 190 return normed_input_values191 192 def __call__(193 self,194 raw_speech: Union[np.ndarray, list[float], list[np.ndarray], list[list[float]]],195 truncation: bool = True,196 pad_to_multiple_of: Optional[int] = None,197 return_tensors: Optional[Union[str, TensorType]] = None,198 return_attention_mask: Optional[bool] = None,199 padding: Optional[str] = "max_length",200 max_length: Optional[int] = None,201 sampling_rate: Optional[int] = None,202 do_normalize: Optional[bool] = None,203 device: Optional[str] = "cpu",204 return_token_timestamps: Optional[bool] = None,205 **kwargs,206 ) -> BatchFeature:207 """208 Main method to featurize and prepare for the model one or several sequence(s). Implementation uses PyTorch for209 the STFT computation if available, otherwise a slower NumPy based one.210 211 Args:212 raw_speech (`np.ndarray`, `list[float]`, `list[np.ndarray]`, `list[list[float]]`):213 The sequence or batch of sequences to be padded. Each sequence can be a numpy array, a list of float214 values, a list of numpy arrays or a list of list of float values. Must be mono channel audio, not215 stereo, i.e. single float per timestep.216 truncation (`bool`, *optional*, default to `True`):217 Activates truncation to cut input sequences longer than *max_length* to *max_length*.218 pad_to_multiple_of (`int`, *optional*, defaults to None):219 If set will pad the sequence to a multiple of the provided value.220 221 This is especially useful to enable the use of Tensor Cores on NVIDIA hardware with compute capability222 `>= 7.5` (Volta), or on TPUs which benefit from having sequence lengths be a multiple of 128.223 return_attention_mask (`bool`, *optional*):224 Whether to return the attention mask. If left to the default, will return the attention mask according225 to the specific feature_extractor's default.226 227 [What are attention masks?](../glossary#attention-mask)228 229 <Tip>230 231 For Whisper models, `attention_mask` should always be passed for batched inference, to avoid subtle232 bugs.233 234 </Tip>235 236 return_tensors (`str` or [`~utils.TensorType`], *optional*):237 If set, will return tensors instead of list of python integers. Acceptable values are:238 239 - `'tf'`: Return TensorFlow `tf.constant` objects.240 - `'pt'`: Return PyTorch `torch.Tensor` objects.241 - `'np'`: Return Numpy `np.ndarray` objects.242 sampling_rate (`int`, *optional*):243 The sampling rate at which the `raw_speech` input was sampled. It is strongly recommended to pass244 `sampling_rate` at the forward call to prevent silent errors and allow automatic speech recognition245 pipeline.246 padding_value (`float`, *optional*, defaults to 0.0):247 The value that is used to fill the padding values / vectors.248 do_normalize (`bool`, *optional*, defaults to `False`):249 Whether or not to zero-mean unit-variance normalize the input. Normalizing can help to significantly250 improve the performance of the model.251 device (`str`, *optional*, defaults to `'cpu'`):252 Specifies the device for computation of the log-mel spectrogram of audio signals in the253 `_torch_extract_fbank_features` method. (e.g., "cpu", "cuda")254 return_token_timestamps (`bool`, *optional*, defaults to `None`):255 Deprecated. Use `return_attention_mask` instead from which the number of frames can be inferred.256 257 Whether or not to return the number of frames of the input raw_speech.258 These num_frames can be used by the model to compute word level timestamps.259 """260 if sampling_rate is not None:261 if sampling_rate != self.sampling_rate:262 raise ValueError(263 f"The model corresponding to this feature extractor: {self.__class__.__name__} was trained using a"264 f" sampling rate of {self.sampling_rate}. Please make sure that the provided `raw_speech` input"265 f" was sampled with {self.sampling_rate} and not {sampling_rate}."266 )267 else:268 logger.warning(269 f"It is strongly recommended to pass the `sampling_rate` argument to `{self.__class__.__name__}()`. "270 "Failing to do so can result in silent errors that might be hard to debug."271 )272 273 is_batched_numpy = isinstance(raw_speech, np.ndarray) and len(raw_speech.shape) > 1274 if is_batched_numpy and len(raw_speech.shape) > 2:275 raise ValueError(f"Only mono-channel audio is supported for input to {self}")276 is_batched = is_batched_numpy or (277 isinstance(raw_speech, (list, tuple)) and (isinstance(raw_speech[0], (np.ndarray, tuple, list)))278 )279 280 if is_batched:281 raw_speech = [np.asarray([speech], dtype=np.float32).T for speech in raw_speech]282 elif not is_batched and not isinstance(raw_speech, np.ndarray):283 raw_speech = np.asarray(raw_speech, dtype=np.float32)284 elif isinstance(raw_speech, np.ndarray) and raw_speech.dtype is np.dtype(np.float64):285 raw_speech = raw_speech.astype(np.float32)286 287 # always return batch288 if not is_batched:289 raw_speech = [np.asarray([raw_speech]).T]290 291 batched_speech = BatchFeature({"input_features": raw_speech})292 293 # convert into correct format for padding294 295 padded_inputs = self.pad(296 batched_speech,297 padding=padding,298 max_length=max_length if max_length else self.n_samples,299 truncation=truncation,300 pad_to_multiple_of=pad_to_multiple_of,301 return_attention_mask=return_attention_mask or do_normalize,302 )303 304 # zero-mean and unit-variance normalization305 if do_normalize:306 padded_inputs["input_features"] = self.zero_mean_unit_var_norm(307 padded_inputs["input_features"],308 attention_mask=padded_inputs["attention_mask"],309 padding_value=self.padding_value,310 )311 padded_inputs["input_features"] = np.stack(padded_inputs["input_features"], axis=0)312 313 # make sure list is in array format314 input_features = padded_inputs.get("input_features").transpose(2, 0, 1)315 316 extract_fbank_features = (317 self._torch_extract_fbank_features if is_torch_available() else self._np_extract_fbank_features318 )319 input_features = extract_fbank_features(input_features[0], device)320 321 if isinstance(input_features[0], list):322 padded_inputs["input_features"] = [np.asarray(feature, dtype=np.float32) for feature in input_features]323 324 else:325 padded_inputs["input_features"] = input_features326 327 if return_attention_mask:328 # rescale from sample (48000) to feature (3000)329 rescaled_attention_mask = padded_inputs["attention_mask"][:, :: self.hop_length]330 331 # The STFT computation produces L//hop_length + 1 frames, but we skip the last frame (see `_torch_extract_fbank_features`).332 # This means we need to trim the rescaled attention mask to match the actual number of frames (L//hop_length) when the input length333 # is not perfectly divisible by the hop length.334 if padded_inputs["attention_mask"].shape[1] % self.hop_length != 0:335 rescaled_attention_mask = rescaled_attention_mask[:, :-1]336 padded_inputs["attention_mask"] = rescaled_attention_mask337 338 if return_token_timestamps is not None:339 logger.warning_once(340 f"`return_token_timestamps` is deprecated for {self.__class__.__name__} and will be removed in Transformers v5. Use `return_attention_mask` instead, as the number of frames can be inferred from it."341 )342 padded_inputs["num_frames"] = [len(raw_speech_i) // self.hop_length for raw_speech_i in raw_speech]343 344 if return_tensors is not None:345 padded_inputs = padded_inputs.convert_to_tensors(return_tensors)346 347 return padded_inputs348 349 350__all__ = ["WhisperFeatureExtractor"]351