mispeech/ced-small

1# coding=utf-82# Copyright 2023 Xiaomi Corporation and 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 CED.17"""18 19from typing import List, Optional, Union20 21import numpy as np22import torch23import torchaudio.transforms as audio_transforms24 25from transformers.feature_extraction_sequence_utils import SequenceFeatureExtractor26from transformers.feature_extraction_utils import BatchFeature27from transformers.utils import logging28 29 30logger = logging.get_logger(__name__)31 32 33class CedFeatureExtractor(SequenceFeatureExtractor):34    r"""35    CedFeatureExtractor extracts Mel spectrogram features from audio signals.36 37    Args:38        f_min (int, *optional*, defaults to 0): Minimum frequency for the Mel filterbank.39        sampling_rate (int, *optional*, defaults to 16000):40            Sampling rate of the input audio signal.41        win_size (int, *optional*, defaults to 512): Window size for the STFT.42        center (bool, *optional*, defaults to `True`):43            Whether to pad the signal on both sides to center it.44        n_fft (int, *optional*, defaults to 512): Number of FFT points for the STFT.45        f_max (int, optional, *optional*): Maximum frequency for the Mel filterbank.46        hop_size (int, *optional*, defaults to 160): Hop size for the STFT.47        feature_size (int, *optional*, defaults to 64): Number of Mel bands to generate.48        padding_value (float, *optional*, defaults to 0.0): Value for padding.49 50    Returns:51        BatchFeature: A BatchFeature object containing the extracted features.52    """53 54    def __init__(55        self,56        f_min: int = 0,57        sampling_rate: int = 16000,58        win_size: int = 512,59        center: bool = True,60        n_fft: int = 512,61        f_max: Optional[int] = None,62        hop_size: int = 160,63        feature_size: int = 64,64        padding_value: float = 0.0,65        **kwargs,66    ):67        super().__init__(68            feature_size=feature_size,69            sampling_rate=sampling_rate,70            padding_value=padding_value,71            **kwargs,72        )73        self.f_min = f_min74        self.win_size = win_size75        self.center = center76        self.n_fft = n_fft77        self.f_max = f_max78        self.hop_size = hop_size79 80        self.model_input_names = ["input_values"]81 82    def __call__(83        self,84        x: Union[np.ndarray, torch.Tensor, List[np.ndarray], List[torch.Tensor]],85        sampling_rate: Optional[int] = None,86        max_length: Optional[int] = None,87        truncation: bool = False,88        return_tensors="pt",89    ) -> BatchFeature:90        r"""91        Extracts Mel spectrogram features from an audio signal tensor.92 93        Args:94            x: Input audio signal tensor.95            sampling_rate (int, *optional*, defaults to `None`):96                Sampling rate of the input audio signal.97            max_length (int, *optional*, defaults to None):98                Maximum length of the input audio signal.99            truncation (bool, *optional*, defaults to `False`):100                Whether to truncate the input signal to max_length.101            return_tensors (str, *optional*, defaults to "pt"):102                If set to "pt", the return type will be a PyTorch tensor.103 104        Returns:105            BatchFeature: A dictionary containing the extracted features.106        """107        if sampling_rate is None:108            sampling_rate = self.sampling_rate109 110        if return_tensors != "pt":111            raise NotImplementedError("Only return_tensors='pt' is currently supported.")112 113        mel_spectrogram = audio_transforms.MelSpectrogram(114            f_min=self.f_min,115            sample_rate=sampling_rate,116            win_length=self.win_size,117            center=self.center,118            n_fft=self.n_fft,119            f_max=self.f_max,120            hop_length=self.hop_size,121            n_mels=self.feature_size,122        )123        amplitude_to_db = audio_transforms.AmplitudeToDB(top_db=120)124 125        if isinstance(x, np.ndarray):126            if x.ndim == 1:127                x = x[np.newaxis, :]128            if x.ndim != 2:129                raise ValueError("np.ndarray input must be a 1D or 2D.")130            x = torch.from_numpy(x)131        elif isinstance(x, torch.Tensor):132            if x.dim() == 1:133                x = x.unsqueeze(0)134            if x.dim() != 2:135                raise ValueError("torch.Tensor input must be a 1D or 2D.")136        elif isinstance(x, (list, tuple)):137            longest_length = max(x_.shape[0] for x_ in x)138            if not truncation and max_length is not None and max_length < longest_length:139                max_length = longest_length140            if not truncation and max_length is None:141                max_length = longest_length142 143 144            if all(isinstance(x_, np.ndarray) for x_ in x):145                if not all(x_.ndim == 1 for x_ in x):146                    raise ValueError("All np.ndarray in a list must be 1D.")147 148                x_trim = [x_[:max_length] for x_ in x]149                x_pad = [np.pad(x_, (0, max_length - x_.shape[0]), mode="constant", constant_values=0) for x_ in x_trim]150                x = torch.stack([torch.from_numpy(x_) for x_ in x_pad])151            elif all(isinstance(x_, torch.Tensor) for x_ in x):152                if not all(x_.dim() == 1 for x_ in x):153                    raise ValueError("All torch.Tensor in a list must be 1D.")154                x_pad = [torch.nn.functional.pad(x_, (0, max_length - x_.shape[0]), value=0) for x_ in x]155                x = torch.stack(x_pad)156            else:157                raise ValueError("Input list must be numpy arrays or PyTorch tensors.")158        else:159            raise ValueError(160                "Input must be a numpy array, a list of numpy arrays, a PyTorch tensor, or a list of PyTorch tensor."161            )162 163        x = x.float()164        x = mel_spectrogram(x)165        x = amplitude_to_db(x)166        return BatchFeature({"input_values": x})167