Add AutoFeatureExtractor support

config.json

@@ -6,7 +6,8 @@
   "auto_map": {
     "AutoConfig": "config.Config",
     "AutoModel": "model.Model",
-    "AutoModelForCTC": "model.Model"
+    "AutoModelForCTC": "model.Model",
+    "AutoFeatureExtractor": "feature_extraction.FeatureExtractor"
   },
   "input_features": 80,
   "vocab_size": 256,

feature_extraction.py

@@ -0,0 +1,378 @@
+"""Feature extraction for ASR model."""
+
+from typing import List, Optional, Union
+
+import numpy as np
+import torch
+from transformers import SequenceFeatureExtractor
+from transformers.audio_utils import mel_filter_bank
+
+
+class FeatureExtractor(SequenceFeatureExtractor):
+    """Feature extractor for ASR model that extracts MFCC features from audio.
+
+    Parameters
+    ----------
+    window_size_ms : int
+        Window size in milliseconds for STFT, default 25.
+    window_stride_ms : int
+        Window stride in milliseconds for STFT, default 10.
+    mel_lower_edge_hertz : int
+        Lower edge of mel frequency range, default 0.
+    mel_upper_edge_hertz : int
+        Upper edge of mel frequency range, default 8000.
+    mel_num_bins : int
+        Number of mel filterbank features, default 80.
+    sample_rate : int
+        Sample rate of audio input, default 16000.
+    padding_value : float
+        Value to use for padding variable-length inputs, default 1000.
+    """
+
+    model_input_names = ["input_features"]
+
+    def __init__(
+        self,
+        window_size_ms: int = 25,
+        window_stride_ms: int = 10,
+        mel_lower_edge_hertz: int = 0,
+        mel_upper_edge_hertz: int = 8000,
+        mel_num_bins: int = 80,
+        sample_rate: int = 16000,
+        padding_value: float = 1000.0,
+        **kwargs,
+    ):
+        super().__init__(
+            feature_size=mel_num_bins,
+            sampling_rate=sample_rate,
+            padding_value=padding_value,
+            **kwargs,
+        )
+
+        self.window_size_ms = window_size_ms
+        self.window_stride_ms = window_stride_ms
+        self.mel_lower_edge_hertz = mel_lower_edge_hertz
+        self.mel_upper_edge_hertz = mel_upper_edge_hertz
+        self.mel_num_bins = mel_num_bins
+        self.log_epsilon = 1e-12
+
+        # Calculate window parameters
+        self.window_size_samples = int(
+            round(self.sampling_rate * self.window_size_ms / 1000.0)
+        )
+        self.window_stride_samples = int(
+            round(self.sampling_rate * self.window_stride_ms / 1000.0)
+        )
+        self.fft_len = self.window_size_samples
+
+        # Precompute mel filterbank matrix
+        self.mel_matrix = mel_filter_bank(
+            num_frequency_bins=self.fft_len // 2 + 1,
+            num_mel_filters=self.mel_num_bins,
+            min_frequency=self.mel_lower_edge_hertz,
+            max_frequency=self.mel_upper_edge_hertz,
+            sampling_rate=self.sampling_rate,
+        )
+
+        # Cache for device-specific mel matrix (avoids repeated conversions)
+        self._mel_matrix_cache = {}  # device -> torch.Tensor
+
+        # Default device for feature extraction
+        self._device = torch.device("cpu")
+
+    @property
+    def device(self):
+        """Get the device for feature extraction."""
+        return self._device
+
+    def to(self, device):
+        """Move feature extractor to a device.
+
+        Parameters
+        ----------
+        device : torch.device or str
+            Device to move to (e.g., 'cuda', 'cpu', torch.device('cuda:0'))
+
+        Returns
+        -------
+        self
+        """
+        self._device = torch.device(device)
+        return self
+
+    def cuda(self, device=None):
+        """Move feature extractor to CUDA device.
+
+        Parameters
+        ----------
+        device : int, optional
+            CUDA device index. If None, uses default CUDA device.
+
+        Returns
+        -------
+        self
+        """
+        if device is None:
+            self._device = torch.device("cuda")
+        else:
+            self._device = torch.device(f"cuda:{device}")
+        return self
+
+    def cpu(self):
+        """Move feature extractor to CPU.
+
+        Returns
+        -------
+        self
+        """
+        self._device = torch.device("cpu")
+        return self
+
+    def to_dict(self):
+        """Serialize to dict, excluding non-serializable attributes."""
+        output = super().to_dict()
+        # Remove non-serializable attributes
+        output.pop("_device", None)
+        output.pop("_mel_matrix_cache", None)
+        return output
+
+    def __call__(
+        self,
+        raw_speech: Union[
+            np.ndarray,
+            torch.Tensor,
+            List[float],
+            List[np.ndarray],
+            List[torch.Tensor],
+            List[List[float]],
+        ],
+        sampling_rate: Optional[int] = None,
+        mask: Optional[Union[np.ndarray, torch.Tensor]] = None,
+        **kwargs,
+    ):
+        """Extract MFCC features from raw audio.
+
+        Parameters
+        ----------
+        raw_speech : np.ndarray or torch.Tensor or List[float] or List[np.ndarray] or List[torch.Tensor] or List[List[float]]
+            The raw audio waveform(s) to extract features from. Can be:
+            - A single waveform as a 1D array/tensor
+            - A batch of waveforms as a 2D array/tensor
+            - A list of waveforms (can be variable length, mask auto-generated)
+        sampling_rate : int, optional
+            Sampling rate of the audio. If provided, must match the feature
+            extractor's sampling_rate.
+        mask : np.ndarray or torch.Tensor, optional
+            Mask for the input audio when input is array/tensor. Should have the same
+            shape as raw_speech. Values should be 1 for real audio and 0 for padding.
+            Not used when raw_speech is a list (mask is auto-generated in that case).
+
+        Returns
+        -------
+        torch.Tensor or dict
+            If no output mask is needed, returns the features tensor directly with
+            shape (batch, time, features). If an output mask is computed, returns a
+            dictionary containing:
+            - input_features: Extracted MFCC features of shape (batch, time, features)
+            - mask: Mask for the features of shape (batch, time)
+        """
+
+        # Validate sampling rate
+        if sampling_rate is not None and sampling_rate != self.sampling_rate:
+            raise ValueError(
+                f"The sampling_rate of the provided audio ({sampling_rate}) "
+                f"doesn't match the feature extractor's sampling_rate ({self.sampling_rate})"
+            )
+
+        input_mask = None
+
+        # Handle tensor/array inputs directly (no padding needed)
+        if isinstance(raw_speech, (torch.Tensor, np.ndarray)):
+            # Ensure input is 2D
+            if raw_speech.ndim == 1:
+                raw_speech = (
+                    raw_speech[np.newaxis, :]
+                    if isinstance(raw_speech, np.ndarray)
+                    else raw_speech.unsqueeze(0)
+                )
+                if mask is not None:
+                    mask = (
+                        mask[np.newaxis, :]
+                        if isinstance(mask, np.ndarray)
+                        else mask.unsqueeze(0)
+                    )
+            elif raw_speech.ndim != 2:
+                raise ValueError(f"Input must be 1D or 2D, got {raw_speech.ndim}D")
+
+            # Convert to torch
+            batched_speech = (
+                raw_speech
+                if isinstance(raw_speech, torch.Tensor)
+                else torch.from_numpy(raw_speech)
+            )
+            # Move to device
+            batched_speech = batched_speech.to(self._device)
+
+            if mask is not None:
+                input_mask = (
+                    mask if isinstance(mask, torch.Tensor) else torch.from_numpy(mask)
+                )
+                # Move to device
+                input_mask = input_mask.to(self._device)
+        else:
+            # Handle list inputs (may need padding)
+            if not isinstance(raw_speech, list):
+                raw_speech = [raw_speech]
+
+            # Convert to torch tensors and move to device
+            torch_speech = []
+            for speech in raw_speech:
+                if isinstance(speech, torch.Tensor):
+                    torch_speech.append(speech.float().to(self._device))
+                else:
+                    torch_speech.append(
+                        torch.from_numpy(np.asarray(speech, dtype=np.float32)).to(
+                            self._device
+                        )
+                    )
+
+            # Find max length and pad to it
+            max_length = max(len(speech) for speech in torch_speech)
+
+            # Pad all sequences to max_length and create mask
+            padded_speech = []
+            masks = []
+            for speech in torch_speech:
+                original_length = len(speech)
+                if original_length < max_length:
+                    padding = torch.full(
+                        (max_length - original_length,),
+                        self.padding_value,
+                        dtype=speech.dtype,
+                        device=self._device,
+                    )
+                    speech = torch.cat([speech, padding])
+
+                    # Create mask: 1 for real data, 0 for padding
+                    mask = torch.ones(max_length, dtype=torch.bool, device=self._device)
+                    mask[original_length:] = 0
+                else:
+                    mask = torch.ones(max_length, dtype=torch.bool, device=self._device)
+
+                padded_speech.append(speech)
+                masks.append(mask)
+
+            # Stack into batch
+            batched_speech = torch.stack(padded_speech, dim=0)
+            input_mask = torch.stack(masks, dim=0)
+
+        # Extract features
+        with torch.no_grad():
+            features = self._extract_features(batched_speech)
+
+        # Compute output mask if we have an input mask
+        output_mask = None
+        if input_mask is not None:
+            output_mask = self._compute_mask(input_mask)
+            # Set masked features to padding_value
+            # output_mask is (batch, time), features is (batch, time, features)
+            # Need to expand mask to broadcast: (batch, time, 1)
+            mask_expanded = output_mask.unsqueeze(-1)
+            features = torch.where(
+                mask_expanded,
+                features,
+                torch.tensor(
+                    self.padding_value, dtype=features.dtype, device=features.device
+                ),
+            )
+
+        # Return features directly if no mask, otherwise return dict
+        if output_mask is not None:
+            return {
+                "input_features": features,
+                "mask": output_mask,
+            }
+        else:
+            return features
+
+    def _extract_features(self, waveform: torch.Tensor) -> torch.Tensor:
+        """Extract MFCC features from waveform.
+
+        Parameters
+        ----------
+        waveform : torch.Tensor
+            Input waveform of shape (batch, time)
+
+        Returns
+        -------
+        torch.Tensor
+            Log mel spectrogram features of shape (batch, time_frames, mel_bins)
+        """
+        # Zero pad if there isn't enough data for at least one frame
+        if waveform.shape[1] < self.window_size_samples:
+            padding = self.window_size_samples - waveform.shape[1]
+            waveform = torch.nn.functional.pad(waveform, (0, padding))
+
+        # Compute spectrogram using STFT
+        spectrogram = torch.stft(
+            waveform,
+            n_fft=self.fft_len,
+            hop_length=self.window_stride_samples,
+            win_length=self.window_size_samples,
+            window=torch.hann_window(self.window_size_samples, device=waveform.device),
+            center=False,
+            return_complex=True,
+        )
+
+        # Take absolute value to get magnitude
+        spectrogram = torch.abs(spectrogram)
+
+        # Get mel matrix from cache or create it
+        device = spectrogram.device
+        dtype = spectrogram.dtype
+        cache_key = (device, dtype)
+
+        if cache_key not in self._mel_matrix_cache:
+            # Convert and cache the mel matrix for this device/dtype combination
+            self._mel_matrix_cache[cache_key] = torch.from_numpy(self.mel_matrix).to(
+                device=device, dtype=dtype
+            )
+
+        mel_matrix = self._mel_matrix_cache[cache_key]
+
+        # Apply mel filterbank: (batch, freq, time) @ (freq, mel) -> (batch, time, mel)
+        # Need to transpose spectrogram from (batch, freq, time) to (batch, time, freq)
+        spectrogram = spectrogram.transpose(1, 2)
+        mel_spectrogram = torch.matmul(spectrogram, mel_matrix)
+
+        # Compute log (with epsilon for stability)
+        log_mel_spectrogram = torch.log(
+            torch.clamp(mel_spectrogram, min=self.log_epsilon)
+        )
+
+        return log_mel_spectrogram
+
+    def _compute_mask(self, input_mask: torch.Tensor) -> torch.Tensor:
+        """Compute output mask for features based on input mask.
+
+        Parameters
+        ----------
+        input_mask : torch.Tensor
+            Input mask of shape (batch, time) with 1 for real data, 0 for padding
+
+        Returns
+        -------
+        torch.Tensor
+            Output mask of shape (batch, time_frames) where a frame is True only if
+            all samples in that frame were valid (not padded)
+        """
+        # Split mask into frames using unfold
+        # unfold(dimension, size, step)
+        mask_frames = input_mask.unfold(
+            1, self.window_size_samples, self.window_stride_samples
+        )
+
+        # A frame is valid only if ALL samples in that frame are valid
+        output_mask = torch.all(mask_frames, dim=-1)
+
+        return output_mask

preprocessor_config.json

@@ -0,0 +1,14 @@
+{
+  "feature_extractor_type": "FeatureExtractor",
+  "processor_class": "FeatureExtractor",
+  "auto_map": {
+    "AutoFeatureExtractor": "feature_extraction.FeatureExtractor"
+  },
+  "window_size_ms": 25,
+  "window_stride_ms": 10,
+  "mel_lower_edge_hertz": 0,
+  "mel_upper_edge_hertz": 8000,
+  "mel_num_bins": 80,
+  "sample_rate": 16000,
+  "padding_value": 1000.0
+}

speech_features.py

@@ -1,125 +0,0 @@
-"""A layer for extracting features from speech data."""
-
-from typing import Iterable, Optional
-
-import keras
-import torch
-from keras import ops
-from numpy.typing import NDArray
-from transformers.audio_utils import mel_filter_bank
-
-
-class SpeechFeatures(keras.layers.Layer):
-    """
-    Computes MFCC features from audio signals.
-    """
-
-    def __init__(
-        self,
-        window_size_ms=25,
-        window_stride_ms=10,
-        mel_lower_edge_hertz=0,
-        mel_upper_edge_hertz=8000,
-        mel_num_bins=80,
-        sample_rate=16000,
-    ):
-        super().__init__()
-
-        self.window_size_ms = window_size_ms
-        self.window_stride_ms = window_stride_ms
-        self.mel_lower_edge_hertz = mel_lower_edge_hertz
-        self.mel_upper_edge_hertz = mel_upper_edge_hertz
-        self.mel_num_bins = mel_num_bins
-        self.sample_rate = sample_rate
-        self.log_epsilon = 1e-12
-
-        self.window_size_samples = int(
-            round(self.sample_rate * self.window_size_ms / 1000.0)
-        )
-        self.window_stride_samples = int(
-            round(self.sample_rate * self.window_stride_ms / 1000.0)
-        )
-
-        self.supports_masking = True
-        self.fft_len = self.window_size_samples
-
-    def build(self, input_shape: Iterable[int]) -> None:
-        # precompute the mel matrix
-        self.mel_matrix = mel_filter_bank(
-            num_frequency_bins=self.fft_len // 2 + 1,
-            num_mel_filters=self.mel_num_bins,
-            min_frequency=self.mel_lower_edge_hertz,
-            max_frequency=self.mel_upper_edge_hertz,
-            sampling_rate=self.sample_rate,
-        )
-
-    def call(self, inputs: NDArray) -> NDArray:
-        """Apply this layer to inputs."""
-
-        if len(inputs.shape) != 2:  # [Batch, Time]
-            raise ValueError(f"Input rank ({len(inputs.shape)}) must be 2")
-
-        # Zero pad if there isn't enough data for at least one frame (so we don't end up
-        # with size 0 axes)
-        inp = ops.pad(
-            inputs,
-            [
-                (0, 0),
-                (
-                    0,
-                    ops.maximum(self.window_size_samples - ops.shape(inputs)[1], 0),
-                ),
-            ],
-        )
-
-        # compute spectrogram
-        spectrogram = self.spectrogram(inp)
-
-        # compute mel spectrogram
-        outputs = self.log_mel(spectrogram)
-
-        return outputs
-
-    def spectrogram(self, inputs: NDArray) -> NDArray:
-        """Compute spectrogram from raw audio."""
-
-        spectrogram = ops.stft(
-            inputs,
-            self.window_size_samples,
-            self.window_stride_samples,
-            fft_length=self.fft_len,
-            center=False,
-        )
-
-        spectrogram = torch.complex(*spectrogram)
-
-        spectrogram = ops.abs(spectrogram)
-
-        return spectrogram
-
-    def log_mel(self, spectrogram: NDArray) -> NDArray:
-        """Transform spectrogram into (log) Mel scale."""
-
-        # multiply spectrogram by mel matrix
-        mel_spectrogram = ops.tensordot(spectrogram, self.mel_matrix, 1)
-
-        # compute log (with epsilon for stability)
-        log_mel_spectrogram = ops.log(ops.maximum(mel_spectrogram, self.log_epsilon))
-
-        return log_mel_spectrogram
-
-    def compute_mask(
-        self, inputs: NDArray, previous_mask: Optional[NDArray] = None
-    ) -> Optional[NDArray]:
-        if previous_mask is None:
-            return None
-
-        # split up mask into frames
-        mask = ops.extract_sequences(
-            previous_mask,
-            self.window_size_samples,
-            self.window_stride_samples,
-        )
-        # mask all the frames that had masked samples in them
-        mask = ops.all(mask, axis=-1)
-        return mask