keras-hub-nightly 0.15.0.dev20240823171555__py3-none-any.whl

keras_hub/src/models/whisper/whisper_audio_feature_extractor.py

@@ -0,0 +1,260 @@
+# Copyright 2024 The KerasHub Authors
+#
+# 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
+#
+#     https://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.
+
+
+import numpy as np
+
+try:
+    import tensorflow as tf
+except ImportError:
+    raise ImportError(
+        "To use `keras_hub`, please install Tensorflow: `pip install tensorflow`. "
+        "The TensorFlow package is required for data preprocessing with any backend."
+    )
+
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.layers.preprocessing.preprocessing_layer import (
+    PreprocessingLayer,
+)
+
+
+@keras_hub_export("keras_hub.models.WhisperAudioFeatureExtractor")
+class WhisperAudioFeatureExtractor(PreprocessingLayer):
+    """
+    Whisper audio feature extractor layer.
+
+    This layer takes in a batch of audio tensors, and computes the log-mel
+    spectrogram features for each audio tensor.
+
+    The input audio tensor can either be of shape `(length_of_audio,)` or
+    `(batch_size, length_of_audio)`. The output is a tensor of shape
+    `(batch_size, num_frames, num_mels)`, where `num_frames` is
+    `(max_audio_length * sampling_rate) / stride`.
+
+    Args:
+        num_mels: int. The number of mel-frequency filters. Defaults to `80`.
+        num_fft_bins: int. The size of the Fourier Transform in STFT.
+            Defaults to `400`.
+        stride: int. The distance between neighboring
+            sliding window frames while computing STFT.
+            Defaults to `160`.
+        sampling_rate: int. The sample rate of the audio. Defaults to `16000`.
+        max_audio_length: int. The length of each audio chunk in
+            seconds. The input audio tensor will be padded/trimmed to
+            `max_audio_length * sampling_rate`. Defaults to `30`.
+
+    Examples:
+
+    ```python
+    audio_tensor = tf.ones((8000,), dtype="float32")
+
+    # Compute the log-mel spectrogram.
+    whisper_audio_feature_extractor = keras_hub.models.WhisperAudioFeatureExtractor()
+    whisper_audio_feature_extractor(audio_tensor)
+
+    # Compute the log-mel spectrogram for a batch of audio tensors.
+    audio_tensor_1 = tf.ones((8000,), dtype="float32")
+    audio_tensor_2 = tf.ones((10000,), dtype="float32"
+    audio_tensor = tf.ragged.stack([audio_tensor_1, audio_tensor_2], axis=0)
+    whisper_audio_feature_extractor(audio_tensor)
+    ```
+    """
+
+    def __init__(
+        self,
+        num_mels=80,
+        num_fft_bins=400,
+        stride=160,
+        sampling_rate=16000,
+        max_audio_length=30,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        self._convert_input_args = False
+        self._allow_non_tensor_positional_args = True
+        self.built = True
+
+        self.num_mels = num_mels
+        self.num_fft_bins = num_fft_bins
+        self.stride = stride
+        self.sampling_rate = sampling_rate
+        self.max_audio_length = max_audio_length
+        self.num_samples = self.sampling_rate * self.max_audio_length
+
+        # After transposition, `self.mel_filters`'s shape is
+        # `(num_fft_bins // 2 + 1, num_mels).`
+        self.mel_filters = self._get_mel_filters()
+
+    def _get_mel_filters(self):
+        """
+        Adapted from Hugging Face
+        (https://github.com/huggingface/transformers/blob/v4.27.1/src/transformers/models/whisper/feature_extraction_whisper.py#L86)
+        """
+
+        # TODO: Convert to TensorFlow ops (if possible).
+
+        dtype = np.float32
+        # Initialize the weights
+        weights = np.zeros(
+            (self.num_mels, int(1 + self.num_fft_bins // 2)), dtype=dtype
+        )
+
+        # Center freqs of each FFT bin
+        fftfreqs = np.fft.rfftfreq(
+            n=self.num_fft_bins, d=1.0 / self.sampling_rate
+        )
+
+        # 'Center freqs' of mel bands - uniformly spaced between limits
+        min_mel = 0.0
+        max_mel = 45.245640471924965
+
+        mels = np.linspace(min_mel, max_mel, self.num_mels + 2)
+
+        mels = np.asanyarray(mels)
+
+        # Fill in the linear scale
+        f_min = 0.0
+        f_sp = 200.0 / 3
+        freqs = f_min + f_sp * mels
+
+        # And now the nonlinear scale
+        min_log_hz = 1000.0  # beginning of log region (Hz)
+        min_log_mel = (min_log_hz - f_min) / f_sp  # same (Mels)
+        logstep = np.log(6.4) / 27.0  # step size for log region
+
+        # If we have vector data, vectorize
+        log_t = mels >= min_log_mel
+        freqs[log_t] = min_log_hz * np.exp(
+            logstep * (mels[log_t] - min_log_mel)
+        )
+
+        mel_f = freqs
+
+        fdiff = np.diff(mel_f)
+        ramps = np.subtract.outer(mel_f, fftfreqs)
+
+        for i in range(self.num_mels):
+            # lower and upper slopes for all bins
+            lower = -ramps[i] / fdiff[i]
+            upper = ramps[i + 2] / fdiff[i + 1]
+
+            # .. then intersect them with each other and zero
+            weights[i] = np.maximum(0, np.minimum(lower, upper))
+
+        # Slaney-style mel is scaled to be approx constant energy per channel
+        enorm = 2.0 / (mel_f[2 : self.num_mels + 2] - mel_f[: self.num_mels])
+        weights *= enorm[:, np.newaxis]
+
+        weights = np.transpose(weights)
+        return tf.constant(weights, dtype=self.compute_dtype)
+
+    def _extract_audio_features(self, audio):
+        audio = tf.cast(audio, self.compute_dtype)
+        # Use "reflection" padding - `tf.signal.stft` uses symmetric padding
+        # internally.
+        audio = tf.pad(
+            audio,
+            paddings=[[0, 0], [self.num_fft_bins // 2, self.num_fft_bins // 2]],
+            mode="REFLECT",
+        )
+
+        # Compute the mel spectrogram.
+        stft = tf.signal.stft(
+            audio,
+            frame_length=self.num_fft_bins,
+            frame_step=self.stride,
+            fft_length=self.num_fft_bins,
+        )
+        magnitudes = tf.square(tf.abs(stft[:, :-1, :]))
+
+        mel_spec = tf.matmul(
+            magnitudes,
+            self.mel_filters,
+        )
+
+        def tf_log10(x):
+            """
+            Computes log base 10 of input tensor using TensorFlow's natural log operator.
+            """
+            numerator = tf.math.log(x)
+            denominator = tf.math.log(tf.constant(10, dtype=numerator.dtype))
+            return numerator / denominator
+
+        # Clamp the values to a minimum value of 1e-10. This is done to avoid
+        # taking the log of 0, i.e., for numerical stability.
+        mel_spec = tf.maximum(mel_spec, 1e-10)
+
+        # Calculate the log mel spectrogram.
+        log_spec = tf_log10(mel_spec)
+        # Dynamic range compression.
+        log_spec_shape = tf.shape(log_spec)
+        max_value_minus_eight = tf.math.subtract(
+            tf.math.reduce_max(log_spec, axis=[1, 2]),
+            tf.cast(8, dtype=log_spec.dtype),
+        )
+        max_value_minus_eight = tf.expand_dims(max_value_minus_eight, axis=1)
+        max_value_minus_eight = tf.repeat(
+            max_value_minus_eight,
+            repeats=log_spec_shape[1] * log_spec_shape[2],
+            axis=1,
+        )
+        max_value_minus_eight = tf.reshape(
+            max_value_minus_eight, shape=log_spec_shape
+        )
+        log_spec = tf.maximum(log_spec, max_value_minus_eight)
+        # Normalization.
+        type_cast_four = tf.cast(4, dtype=log_spec.dtype)
+        log_spec = tf.math.divide(
+            tf.math.add(log_spec, type_cast_four),
+            type_cast_four,
+        )
+
+        return log_spec
+
+    def call(self, audio):
+        if not isinstance(audio, (tf.Tensor, tf.RaggedTensor)):
+            audio = tf.convert_to_tensor(audio)
+
+        rank_1_input = audio.shape.rank == 1
+        if rank_1_input:
+            audio = tf.expand_dims(audio, 0)
+
+        # Convert the tensor to a Ragged Tensor.
+        if isinstance(audio, tf.Tensor):
+            audio = tf.RaggedTensor.from_tensor(audio)
+
+        # Pad audio.
+        audio_shape = audio.shape.as_list()
+        audio_shape[-1] = self.num_samples
+        audio = audio.to_tensor(shape=audio_shape)
+
+        # Find the log mel spectrogram.
+        log_spec = self._extract_audio_features(audio)
+        if rank_1_input:
+            log_spec = tf.squeeze(log_spec, 0)
+        return log_spec
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "num_mels": self.num_mels,
+                "num_fft_bins": self.num_fft_bins,
+                "stride": self.stride,
+                "sampling_rate": self.sampling_rate,
+                "max_audio_length": self.max_audio_length,
+            }
+        )
+        return config

keras_hub/src/models/whisper/whisper_backbone.py

@@ -0,0 +1,305 @@
+# Copyright 2024 The KerasHub Authors
+#
+# 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
+#
+#     https://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.
+
+
+import keras
+from keras import ops
+
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.layers.modeling.position_embedding import PositionEmbedding
+from keras_hub.src.layers.modeling.token_and_position_embedding import (
+    TokenAndPositionEmbedding,
+)
+from keras_hub.src.models.backbone import Backbone
+from keras_hub.src.models.whisper.whisper_decoder import WhisperDecoder
+from keras_hub.src.models.whisper.whisper_encoder import WhisperEncoder
+from keras_hub.src.utils.tensor_utils import assert_tf_backend
+
+
+def whisper_kernel_initializer(stddev=0.02):
+    return keras.initializers.TruncatedNormal(stddev=stddev)
+
+
+class Padder(keras.layers.Layer):
+    def call(self, x):
+        return ops.pad(x, [[0, 0], [1, 1], [0, 0]])
+
+
+@keras_hub_export("keras_hub.models.WhisperBackbone")
+class WhisperBackbone(Backbone):
+    """A Whisper encoder-decoder network for speech.
+
+    This class implements a Transformer-based encoder-decoder model as
+    described in
+    ["Robust Speech Recognition via Large-Scale Weak Supervision"](https://arxiv.org/abs/2212.04356).
+    It includes the embedding lookups and transformer layers, but not the head
+    for predicting the next token.
+
+    The default constructor gives a fully customizable, randomly initialized Whisper
+    model with any number of layers, heads, and embedding dimensions. To load
+    preset architectures and weights, use the `from_preset()` constructor.
+
+    Disclaimer: Pre-trained models are provided on an "as is" basis, without
+    warranties or conditions of any kind. The underlying model is provided by a
+    third party and subject to a separate license, available
+    [here](https://github.com/openai/whisper).
+
+    Args:
+        vocabulary_size: int. The size of the token vocabulary.
+        num_layers: int. The number of transformer encoder layers and
+            transformer decoder layers.
+        num_heads: int. The number of attention heads for each transformer.
+            The hidden size must be divisible by the number of attention heads.
+        hidden_dim: int. The size of the transformer encoding and pooler layers.
+        intermediate_dim: int. The output dimension of the first Dense layer in
+            a two-layer feedforward network for each transformer.
+        num_mels: int. The number of mel-frequency filters. Defaults to `80`.
+        dropout: float. Dropout probability for the Transformer encoder.
+        max_encoder_sequence_length: int. The maximum sequence length that the
+            audio encoder can consume. Since the second convolutional layer in
+            the encoder reduces the sequence length by half (stride of 2), we
+            use `max_encoder_sequence_length // 2` as the sequence length for the
+            positional embedding layer.
+        max_decoder_sequence_length: int. The maximum sequence length that the
+            text decoder can consume.
+        dtype: string or `keras.mixed_precision.DTypePolicy`. The dtype to use
+            for model computations and weights. Note that some computations,
+            such as softmax and layer normalization, will always be done at
+            float32 precision regardless of dtype.
+
+    Examples:
+
+    ```python
+    input_data = {
+        "encoder_features": np.ones(shape=(1, 12, 80), dtype="int32"),
+        "decoder_token_ids": np.ones(shape=(1, 12), dtype="int32"),
+        "decoder_padding_mask": np.array(
+            [[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0]]
+        ),
+    }
+
+    # Randomly initialized Whisper encoder-decoder model with a custom config.
+    model = keras_hub.models.WhisperBackbone(
+        vocabulary_size=51864,
+        num_layers=4,
+        num_heads=4,
+        hidden_dim=256,
+        intermediate_dim=512,
+        max_encoder_sequence_length=128,
+        max_decoder_sequence_length=128,
+    )
+    model(input_data)
+    ```
+    """
+
+    def __init__(
+        self,
+        vocabulary_size,
+        num_layers,
+        num_heads,
+        hidden_dim,
+        intermediate_dim,
+        num_mels=80,
+        dropout=0.0,
+        max_encoder_sequence_length=3000,
+        max_decoder_sequence_length=448,
+        dtype=None,
+        **kwargs,
+    ):
+        assert_tf_backend(self.__class__.__name__)
+
+        # === Layers ===
+        self.encoder_conv_layer_1 = keras.layers.Conv1D(
+            filters=hidden_dim,
+            kernel_size=3,
+            strides=1,
+            padding="same",
+            dtype=dtype,
+            name="encoder_token_embedding_conv_layer_1",
+        )
+        self.encoder_conv_layer_2 = keras.layers.Conv1D(
+            filters=hidden_dim,
+            kernel_size=3,
+            strides=2,
+            padding="valid",
+            dtype=dtype,
+            name="encoder_token_embedding_conv_layer_2",
+        )
+        self.encoder_padder = Padder(
+            dtype=dtype,
+            name="encoder_padder",
+        )
+        self.encoder_position_embedding = PositionEmbedding(
+            initializer=whisper_kernel_initializer(),
+            sequence_length=max_encoder_sequence_length // 2,
+            dtype=dtype,
+            name="encoder_position_embedding",
+            trainable=False,
+        )
+        self.encoder_embeddings_add = keras.layers.Add(
+            dtype=dtype,
+            name="encoder_embeddings_add",
+        )
+        self.encoder_embeddings_dropout = keras.layers.Dropout(
+            dropout,
+            dtype=dtype,
+            name="encoder_embeddings_dropout",
+        )
+        self.encoder_transformer_layers = []
+        for i in range(num_layers):
+            layer = WhisperEncoder(
+                num_heads=num_heads,
+                intermediate_dim=intermediate_dim,
+                activation=keras.activations.gelu,
+                layer_norm_epsilon=1e-5,
+                dropout=dropout,
+                kernel_initializer=whisper_kernel_initializer(),
+                normalize_first=True,
+                dtype=dtype,
+                name=f"transformer_encoder_layer_{i}",
+            )
+            self.encoder_transformer_layers.append(layer)
+        self.encoder_layer_norm = keras.layers.LayerNormalization(
+            axis=-1,
+            epsilon=1e-5,
+            dtype=dtype,
+            name="encoder_layer_norm",
+        )
+        self.decoder_embeddings = TokenAndPositionEmbedding(
+            vocabulary_size=vocabulary_size,
+            sequence_length=max_decoder_sequence_length,
+            embedding_dim=hidden_dim,
+            embeddings_initializer=whisper_kernel_initializer(),
+            dtype=dtype,
+            name="decoder_token_and_position_embedding",
+        )
+        self.token_embedding = self.decoder_embeddings.token_embedding
+        self.decoder_embeddings_dropout = keras.layers.Dropout(
+            dropout,
+            dtype=dtype,
+            name="decoder_embeddings_dropout",
+        )
+        self.decoder_transformer_layers = []
+        for i in range(num_layers):
+            layer = WhisperDecoder(
+                intermediate_dim=intermediate_dim,
+                num_heads=num_heads,
+                dropout=dropout,
+                activation=keras.activations.gelu,
+                layer_norm_epsilon=1e-5,
+                kernel_initializer=whisper_kernel_initializer(),
+                normalize_first=True,
+                dtype=dtype,
+                name=f"transformer_decoder_layer_{i}",
+            )
+            self.decoder_transformer_layers.append(layer)
+        self.decoder_layer_norm = keras.layers.LayerNormalization(
+            axis=-1,
+            epsilon=1e-5,
+            dtype=dtype,
+            name="decoder_layer_norm",
+        )
+
+        # === Functional Model ===
+        # Note that the encoder does not have a padding mask:
+        # https://github.com/openai/whisper/blob/v20230124/whisper/model.py#L132.
+        encoder_feature_input = keras.Input(
+            shape=(None, num_mels), dtype="float32", name="encoder_features"
+        )
+        decoder_token_id_input = keras.Input(
+            shape=(None,), dtype="int32", name="decoder_token_ids"
+        )
+        decoder_padding_mask_input = keras.Input(
+            shape=(None,), dtype="int32", name="decoder_padding_mask"
+        )
+        # Encoder.
+        # Embed the input features. This consists of two 1D convolutional
+        # layers.
+        # For the first layer, we use `padding="same"` since that corresponds to
+        # a padding size of 1.
+        embedded_features = keras.activations.gelu(
+            self.encoder_conv_layer_1(encoder_feature_input),
+            approximate=False,
+        )
+        # For the second conv. layer, we cannot use `padding="same"` since
+        # that corresponds to a padding size of 1.5 (since stride is 2). Hence,
+        # we will manually pad the input.
+        embedded_features = self.encoder_padder(embedded_features)
+        embedded_features = keras.activations.gelu(
+            self.encoder_conv_layer_2(embedded_features),
+            approximate=False,
+        )
+        # The position embedding layer for the encoder is a sinusoidal embedding
+        # layer: https://github.com/openai/whisper/blob/v20230124/whisper/model.py#L137.
+        # Hence, we set it to be non-trainable.
+        # TODO: We can use `keras_hub.layers.SinePositionEncoding` layer.
+        positions = self.encoder_position_embedding(embedded_features)
+        x = self.encoder_embeddings_add((embedded_features, positions))
+        x = self.encoder_embeddings_dropout(x)
+        for transformer_layer in self.encoder_transformer_layers:
+            x = transformer_layer(x)
+        x = self.encoder_layer_norm(x)
+        encoder_output = x
+        # Decoder.
+        x = self.decoder_embeddings(decoder_token_id_input)
+        x = self.decoder_embeddings_dropout(x)
+        for transformer_layer in self.decoder_transformer_layers:
+            x = transformer_layer(
+                decoder_sequence=x,
+                encoder_sequence=encoder_output,
+                decoder_padding_mask=decoder_padding_mask_input,
+            )
+        x = self.decoder_layer_norm(x)
+        decoder_output = x
+        super().__init__(
+            inputs={
+                "encoder_features": encoder_feature_input,
+                "decoder_token_ids": decoder_token_id_input,
+                "decoder_padding_mask": decoder_padding_mask_input,
+            },
+            outputs={
+                "encoder_sequence_output": encoder_output,
+                "decoder_sequence_output": decoder_output,
+            },
+            dtype=dtype,
+            **kwargs,
+        )
+
+        # === Config ===
+        self.vocabulary_size = vocabulary_size
+        self.num_layers = num_layers
+        self.num_heads = num_heads
+        self.hidden_dim = hidden_dim
+        self.intermediate_dim = intermediate_dim
+        self.num_mels = num_mels
+        self.dropout = dropout
+        self.max_encoder_sequence_length = max_encoder_sequence_length
+        self.max_decoder_sequence_length = max_decoder_sequence_length
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "vocabulary_size": self.vocabulary_size,
+                "num_layers": self.num_layers,
+                "num_heads": self.num_heads,
+                "hidden_dim": self.hidden_dim,
+                "intermediate_dim": self.intermediate_dim,
+                "num_mels": self.num_mels,
+                "dropout": self.dropout,
+                "max_encoder_sequence_length": self.max_encoder_sequence_length,
+                "max_decoder_sequence_length": self.max_decoder_sequence_length,
+            }
+        )
+        return config