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

keras_hub/src/models/whisper/whisper_preprocessor.py

@@ -0,0 +1,326 @@
+# 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 absl import logging
+
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.layers.preprocessing.start_end_packer import StartEndPacker
+from keras_hub.src.models.preprocessor import Preprocessor
+from keras_hub.src.models.whisper.whisper_audio_feature_extractor import (
+    WhisperAudioFeatureExtractor,
+)
+from keras_hub.src.models.whisper.whisper_tokenizer import WhisperTokenizer
+from keras_hub.src.utils.keras_utils import (
+    convert_inputs_to_list_of_tensor_segments,
+)
+
+
+@keras_hub_export("keras_hub.models.WhisperPreprocessor")
+class WhisperPreprocessor(Preprocessor):
+    """A Whisper preprocessing layer which handles audio and text input.
+
+    This preprocessing layer will do three things:
+
+     1. Compute the log-mel spectrogram of the audio tensor inputs using
+        `audio_feature_extractor`.
+     2. Tokenize decoder inputs using the `tokenizer`.
+     2. Add the appropriate special tokens - `"<|startoftranscript|>", task
+        token, language token, `"<|endoftext|>"`, etc.
+     3. Construct a dictionary with keys `"encoder_features"`,
+        `"decoder_token_ids"`, `"decoder_padding_mask"` that can be passed
+        directly to a Whisper model.
+
+    Args:
+        tokenizer: A `keras_hub.models.WhisperTokenizer` instance.
+        audio_feature_extractor: A
+            `keras_hub.models.WhisperAudioFeatureExtractor` instance or `None`.
+            If `None` a feature extractor with default parameters will be
+            created.
+        decoder_sequence_length: The length of the packed decoder inputs.
+        language: string, language token. Should only be passed if your
+            tokenizer is multilingual.
+        task: string, task name. One of `"transcribe"`, `"translate"`. Should
+            only be passed if your tokenizer is multilingual.
+        no_timestamps: bool. If True, `"<|no_timestamps|>"` will be added as a
+            special token to your input.
+
+    Call arguments:
+        x: A dictionary with `"encoder_audio"` and `"decoder_text"` as its keys.
+            `"encoder_audio"` should correspond to the input audio tensor.
+            `"decoder_text"` should be a tensor of single string sequences.
+            Inputs may be batched or unbatched. Raw python inputs will be
+            converted to tensors.
+        y: Any label data. Will be passed through unaltered.
+        sample_weight: Any label weight data. Will be passed through unaltered.
+
+    Examples:
+
+    Directly calling the layer on data.
+    ```python
+    preprocessor = keras_hub.models.WhisperPreprocessor.from_preset(
+        "whisper_tiny_en",
+    )
+
+    # Preprocess unbatched inputs.
+    input_data = {
+        "encoder_audio": tf.ones((200,)),
+        "decoder_text": "The quick brown fox jumped.",
+    }
+    preprocessor(input_data)
+
+    # Preprocess batched inputs.
+    input_data = {
+        "encoder_audio": tf.ones((2, 200)),
+        "decoder_text": ["The quick brown fox jumped.", "Call me Ishmael."],
+    }
+    preprocessor(input_data)
+
+    # Custom audio feature extractor and vocabulary.
+    audio_feature_extractor = keras_hub.models.WhisperAudioFeatureExtractor(
+        num_mels=80,
+        num_fft_bins=400,
+        stride=100,
+        sampling_rate=100,
+        max_audio_length=5,
+    )
+
+    features = ["a quick fox.", "a fox quick."]
+    vocab = {"<|endoftext|>": 0, "a": 4, "Ġquick": 5, "Ġfox": 6}
+    merges = ["Ġ q", "u i", "c k", "ui ck", "Ġq uick"]
+    merges += ["Ġ f", "o x", "Ġf ox"]
+    special_tokens = {
+        "<|startoftranscript|>": 9,
+        "<|endoftext|>": 10,
+        "<|notimestamps|>": 11,
+        "<|transcribe|>": 12,
+        "<|translate|>": 13,
+    }
+
+    tokenizer = keras_hub.models.WhisperTokenizer(
+        vocabulary=vocab,
+        merges=merges,
+        special_tokens=special_tokens,
+    )
+    preprocessor = keras_hub.models.WhisperPreprocessor(
+        audio_feature_extractor=audio_feature_extractor,
+        tokenizer=tokenizer,
+    )
+
+    input_data = {
+        "encoder_audio": tf.ones((200,)),
+        "decoder_text": "The quick brown fox jumped.",
+    }
+    preprocessor(input_data)
+    ```
+
+    Mapping with `tf.data.Dataset`.
+    ```python
+    preprocessor = keras_hub.models.WhisperPreprocessor.from_preset(
+        "whisper_tiny_en")
+
+    # Map labeled single sentences.
+    features = {
+        "encoder_audio": tf.ones((2, 200)),
+        "decoder_text": ["The quick brown fox jumped.", "Call me Ishmael."],
+    }
+    labels = tf.constant(["True", "False"])
+    ds = tf.data.Dataset.from_tensor_slices((features, labels))
+    ds = ds.map(preprocessor, num_parallel_calls=tf.data.AUTOTUNE)
+
+    # Map unlabeled single sentences.
+    features = {
+        "encoder_audio": tf.ones((2, 200)),
+        "decoder_text": ["The quick brown fox jumped.", "Call me Ishmael."],
+    }
+    ds = tf.data.Dataset.from_tensor_slices(features)
+    ds = ds.map(preprocessor, num_parallel_calls=tf.data.AUTOTUNE)
+    ```
+    """
+
+    tokenizer_cls = WhisperTokenizer
+
+    def __init__(
+        self,
+        tokenizer,
+        audio_feature_extractor=None,
+        decoder_sequence_length=448,
+        language=None,
+        task=None,
+        no_timestamps=True,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        if audio_feature_extractor is None:
+            audio_feature_extractor = WhisperAudioFeatureExtractor()
+        self.audio_feature_extractor = audio_feature_extractor
+        self.tokenizer = tokenizer
+        self.decoder_packer = None
+        self.decoder_sequence_length = decoder_sequence_length
+        self.language = language
+        self.task = task
+        self.no_timestamps = no_timestamps
+
+    def build(self, input_shape):
+        # Defer packer creation to `build()` so that we can be sure tokenizer
+        # assets have loaded when restoring a saved model.
+
+        # Create list of tokens to be prepended to decoder inputs.
+        bos_tokens = [self.tokenizer.bos_token_id]
+        if self.tokenizer.language_tokens is not None:
+            if (
+                self.language is None
+                or self.language not in self.tokenizer.language_tokens
+            ):
+                raise ValueError(
+                    "You must pass a non-None value for `language` when using "
+                    "a multilingual tokenizer. The value must be one of "
+                    f'{",".join(self.tokenizer.language_tokens.keys())}. '
+                    f"Received: language={self.language}."
+                )
+            if self.task is None or self.task not in [
+                "transcribe",
+                "translate",
+            ]:
+                raise ValueError(
+                    "You must pass a non-None value for `task` when using "
+                    "a multilingual tokenizer. The value must be one of "
+                    '`"transcribe"`, `"translate"`. '
+                    f"Received: task={self.task}."
+                )
+
+            bos_tokens += [self.tokenizer.language_tokens[self.language]]
+
+            if self.task == "transcribe":
+                bos_tokens += [self.tokenizer.special_tokens["<|transcribe|>"]]
+            elif self.task == "translate":
+                bos_tokens += [self.tokenizer.special_tokens["<|translate|>"]]
+        else:
+            if self.language is not None:
+                logging.info(
+                    "`tokenizer` is monolingual, and `language` has a "
+                    "non-`None` value. Setting `language` to `None`."
+                )
+                self.language = None
+            if self.task is not None:
+                logging.info(
+                    "`tokenizer` is monolingual, and `task` has a "
+                    "non-`None` value. Setting `task` to `None`."
+                )
+                self.task = None
+
+        if self.no_timestamps:
+            bos_tokens += [self.tokenizer.no_timestamps_token_id]
+
+        # TODO: Use `MultiSegmentPacker` instead of `StartEndPacker` once we
+        # want to move to multi-segment packing and have improved
+        # `MultiSegmentPacker`'s performance.
+        self.decoder_packer = StartEndPacker(
+            start_value=bos_tokens,
+            end_value=self.tokenizer.eos_token_id,
+            pad_value=self.tokenizer.pad_token_id,
+            sequence_length=self.decoder_sequence_length,
+            return_padding_mask=True,
+        )
+
+    def call(self, x, y=None, sample_weight=None, decoder_sequence_length=None):
+        if not (
+            isinstance(x, dict)
+            and ["encoder_audio", "decoder_text"] == list(x.keys())
+        ):
+            raise ValueError(
+                '`x` must be a dictionary, containing the keys `"encoder_audio"`'
+                f' and `"decoder_text"`. Received x={x}.'
+            )
+
+        encoder_audio = x["encoder_audio"]
+        decoder_text = x["decoder_text"]
+
+        encoder_audio = convert_inputs_to_list_of_tensor_segments(encoder_audio)
+        decoder_text = convert_inputs_to_list_of_tensor_segments(decoder_text)
+
+        if len(encoder_audio) > 1 or len(decoder_text) > 1:
+            raise ValueError(
+                '`WhisperPreprocessor` requires both `"encoder_audio"` and '
+                f'`"decoder_text"` to contain only one segment, but received '
+                f"{len(encoder_audio)} and {len(decoder_text)}, respectively."
+            )
+
+        encoder_features = self.audio_feature_extractor(encoder_audio[0])
+        decoder_sequence_length = (
+            decoder_sequence_length or self.decoder_sequence_length
+        )
+        decoder_inputs = self.tokenizer(decoder_text[0])
+        decoder_token_ids, decoder_padding_mask = self.decoder_packer(
+            decoder_inputs,
+            sequence_length=decoder_sequence_length,
+        )
+
+        x = {
+            "encoder_features": encoder_features,
+            "decoder_token_ids": decoder_token_ids,
+            "decoder_padding_mask": decoder_padding_mask,
+        }
+
+        return keras.utils.pack_x_y_sample_weight(x, y, sample_weight)
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "audio_feature_extractor": keras.layers.serialize(
+                    self.audio_feature_extractor
+                ),
+                "decoder_sequence_length": self.decoder_sequence_length,
+                "language": self.language,
+                "task": self.task,
+                "no_timestamps": self.no_timestamps,
+            }
+        )
+        return config
+
+    @classmethod
+    def from_config(cls, config):
+        if "tokenizer" in config and isinstance(config["tokenizer"], dict):
+            config["tokenizer"] = keras.layers.deserialize(config["tokenizer"])
+
+        if "audio_feature_extractor" in config and isinstance(
+            config["audio_feature_extractor"], dict
+        ):
+            config["audio_feature_extractor"] = keras.layers.deserialize(
+                config["audio_feature_extractor"]
+            )
+
+        return cls(**config)
+
+    @property
+    def decoder_sequence_length(self):
+        """The padded length of decoder input sequences."""
+        return self._decoder_sequence_length
+
+    @decoder_sequence_length.setter
+    def decoder_sequence_length(self, value):
+        self._decoder_sequence_length = value
+        if self.decoder_packer is not None:
+            self.decoder_packer.sequence_length = value
+
+    @property
+    def sequence_length(self):
+        """Alias for `decoder_sequence_length`."""
+        return self.decoder_sequence_length
+
+    @sequence_length.setter
+    def sequence_length(self, value):
+        self.decoder_sequence_length = value

keras_hub/src/models/whisper/whisper_presets.py

@@ -0,0 +1,148 @@
+# 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.
+
+# Metadata for loading pretrained model weights.
+backbone_presets = {
+    "whisper_tiny_en": {
+        "metadata": {
+            "description": (
+                "4-layer Whisper model. Trained on 438,000 hours of labelled "
+                "English speech data."
+            ),
+            "params": 37184256,
+            "official_name": "Whisper",
+            "path": "whisper",
+            "model_card": "https://github.com/openai/whisper/blob/main/model-card.md",
+        },
+        "kaggle_handle": "kaggle://keras/whisper/keras/whisper_tiny_en/2",
+    },
+    "whisper_base_en": {
+        "metadata": {
+            "description": (
+                "6-layer Whisper model. Trained on 438,000 hours of labelled "
+                "English speech data."
+            ),
+            "params": 124439808,
+            "official_name": "Whisper",
+            "path": "whisper",
+            "model_card": "https://github.com/openai/whisper/blob/main/model-card.md",
+        },
+        "kaggle_handle": "kaggle://keras/whisper/keras/whisper_base_en/2",
+    },
+    "whisper_small_en": {
+        "metadata": {
+            "description": (
+                "12-layer Whisper model. Trained on 438,000 hours of labelled "
+                "English speech data."
+            ),
+            "params": 241734144,
+            "official_name": "Whisper",
+            "path": "whisper",
+            "model_card": "https://github.com/openai/whisper/blob/main/model-card.md",
+        },
+        "kaggle_handle": "kaggle://keras/whisper/keras/whisper_small_en/2",
+    },
+    "whisper_medium_en": {
+        "metadata": {
+            "description": (
+                "24-layer Whisper model. Trained on 438,000 hours of labelled "
+                "English speech data."
+            ),
+            "params": 763856896,
+            "official_name": "Whisper",
+            "path": "whisper",
+            "model_card": "https://github.com/openai/whisper/blob/main/model-card.md",
+        },
+        "kaggle_handle": "kaggle://keras/whisper/keras/whisper_medium_en/2",
+    },
+    "whisper_tiny_multi": {
+        "metadata": {
+            "description": (
+                "4-layer Whisper model. Trained on 680,000 hours of labelled "
+                "multilingual speech data."
+            ),
+            "params": 37760640,
+            "official_name": "Whisper",
+            "path": "whisper",
+            "model_card": "https://github.com/openai/whisper/blob/main/model-card.md",
+        },
+        "kaggle_handle": "kaggle://keras/whisper/keras/whisper_tiny_multi/2",
+    },
+    "whisper_base_multi": {
+        "metadata": {
+            "description": (
+                "6-layer Whisper model. Trained on 680,000 hours of labelled "
+                "multilingual speech data."
+            ),
+            "params": 72593920,
+            "official_name": "Whisper",
+            "path": "whisper",
+            "model_card": "https://github.com/openai/whisper/blob/main/model-card.md",
+        },
+        "kaggle_handle": "kaggle://keras/whisper/keras/whisper_base_multi/2",
+    },
+    "whisper_small_multi": {
+        "metadata": {
+            "description": (
+                "12-layer Whisper model. Trained on 680,000 hours of labelled "
+                "multilingual speech data."
+            ),
+            "params": 241734912,
+            "official_name": "Whisper",
+            "path": "whisper",
+            "model_card": "https://github.com/openai/whisper/blob/main/model-card.md",
+        },
+        "kaggle_handle": "kaggle://keras/whisper/keras/whisper_small_multi/2",
+    },
+    "whisper_medium_multi": {
+        "metadata": {
+            "description": (
+                "24-layer Whisper model. Trained on 680,000 hours of labelled "
+                "multilingual speech data."
+            ),
+            "params": 763857920,
+            "official_name": "Whisper",
+            "path": "whisper",
+            "model_card": "https://github.com/openai/whisper/blob/main/model-card.md",
+        },
+        "kaggle_handle": "kaggle://keras/whisper/keras/whisper_medium_multi/2",
+    },
+    "whisper_large_multi": {
+        "metadata": {
+            "description": (
+                "32-layer Whisper model. Trained on 680,000 hours of labelled "
+                "multilingual speech data."
+            ),
+            "params": 1543304960,
+            "official_name": "Whisper",
+            "path": "whisper",
+            "model_card": "https://github.com/openai/whisper/blob/main/model-card.md",
+        },
+        "kaggle_handle": "kaggle://keras/whisper/keras/whisper_large_multi/2",
+    },
+    "whisper_large_multi_v2": {
+        "metadata": {
+            "description": (
+                "32-layer Whisper model. Trained for 2.5 epochs on 680,000  "
+                "hours of labelled multilingual speech data. An improved "
+                "of `whisper_large_multi`."
+            ),
+            "params": 1543304960,
+            "official_name": "Whisper",
+            "path": "whisper",
+            "model_card": "https://github.com/openai/whisper/blob/main/model-card.md",
+        },
+        "kaggle_handle": "kaggle://keras/whisper/keras/whisper_large_multi_v2/2",
+    },
+}

keras_hub/src/models/whisper/whisper_tokenizer.py

@@ -0,0 +1,163 @@
+# 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 json
+
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.tokenizers.byte_pair_tokenizer import BytePairTokenizer
+
+
+def _load_dict(dict_or_path):
+    if isinstance(dict_or_path, str):
+        with open(dict_or_path, "r", encoding="utf-8") as f:
+            dict_or_path = json.load(f)
+    return dict_or_path
+
+
+@keras_hub_export("keras_hub.models.WhisperTokenizer")
+class WhisperTokenizer(BytePairTokenizer):
+    """Whisper text tokenizer using Byte-Pair Encoding subword segmentation.
+
+    This tokenizer class will tokenize raw strings into integer sequences and
+    is based on `keras_hub.tokenizers.BytePairTokenizer`.
+    This tokenizer does not provide truncation or padding of inputs.
+
+    Args:
+        vocabulary: string or dict, maps token to integer ids. If it is a
+            string, it should be the file path to a json file.
+        merges: string or list, contains the merge rule. If it is a string,
+            it should be the file path to merge rules. The merge rule file
+            should have one merge rule per line. Every merge rule contains
+            merge entities separated by a space.
+        special_tokens: string or dict, maps special tokens to integer IDs. If
+            it is a string, it should be the path to a JSON file.
+        language_tokens: string or dict, maps language tokens to integer IDs. If
+            not None, the tokenizer will be assumed to be a multilingual
+            tokenizer.
+    """
+
+    def __init__(
+        self,
+        vocabulary=None,
+        merges=None,
+        special_tokens=None,
+        language_tokens=None,
+        **kwargs,
+    ):
+        special_tokens = _load_dict(special_tokens)
+        if language_tokens is not None:
+            language_tokens = _load_dict(language_tokens)
+
+        # Necessary special tokens.
+        self.bos_token = "<|startoftranscript|>"
+        self.eos_token = "<|endoftext|>"
+        # TODO: The pad token for the multilingual tokenizer is actually
+        # "", but it errors out (OOM). After BPE is fixed, we can update
+        # this to "". For now, we will use `"<|endoftext|>"`.
+        self.pad_token = "<|endoftext|>"
+
+        self.no_timestamps_token = "<|notimestamps|>"
+        # Task special tokens.
+        self.translate_token = "<|translate|>"
+        self.transcribe_token = "<|transcribe|>"
+
+        for token in [
+            self.bos_token,
+            self.eos_token,
+            self.pad_token,
+            self.no_timestamps_token,
+            self.translate_token,
+            self.transcribe_token,
+        ]:
+            if token not in special_tokens:
+                raise ValueError(
+                    f"Cannot find token `'{token}'` in the provided "
+                    f"`special_tokens`. Please provide `'{token}'` in your "
+                    "`special_tokens`."
+                )
+
+        self.bos_token_id = special_tokens[self.bos_token]
+        self.eos_token_id = special_tokens[self.eos_token]
+        self.pad_token_id = special_tokens[self.pad_token]
+        self.no_timestamps_token_id = special_tokens[self.no_timestamps_token]
+        self.translate_token_id = special_tokens[self.translate_token]
+        self.transcribe_token_id = special_tokens[self.transcribe_token]
+
+        self.special_tokens = special_tokens
+        self.language_tokens = language_tokens
+
+        # TODO: Add language tokens to `unsplittable_tokens` once we figure
+        # out the performance issue with a large list.
+        unsplittable_tokens = list(special_tokens.keys())
+
+        super().__init__(
+            vocabulary=vocabulary,
+            merges=merges,
+            unsplittable_tokens=unsplittable_tokens,
+            **kwargs,
+        )
+
+    def save_assets(self, dir_path):
+        # TODO: whisper is currently mutating it's vocabulary before passing
+        # it to the super class, so we need to restore the unmutated vocabulary
+        # before saving our assets. We should find a more robust (and memory
+        # efficient) way to do this.
+        vocabulary = self.vocabulary
+        self.vocabulary = self._initial_vocabulary
+        super().save_assets(dir_path)
+        self.vocabulary = vocabulary
+
+    def set_vocabulary_and_merges(self, vocabulary, merges):
+        if vocabulary is not None:
+            vocabulary = _load_dict(vocabulary)
+            self._initial_vocabulary = dict(vocabulary)
+
+            if self.language_tokens is not None:
+                # Multilingual tokenizer.
+                # Add language tokens to the vocabulary. This makes
+                # detokenization easier for us.
+                vocabulary = {
+                    **vocabulary,
+                    **self.language_tokens,
+                }
+
+            for token in [
+                self.bos_token,
+                self.eos_token,
+                self.pad_token,
+                self.no_timestamps_token,
+                self.translate_token,
+                self.transcribe_token,
+            ]:
+                vocabulary[token] = self.special_tokens[token]
+        else:
+            self._initial_vocabulary = None
+
+        super().set_vocabulary_and_merges(vocabulary, merges)
+
+    def get_config(self):
+        config = super().get_config()
+
+        # In the constructor, we pass the list of special tokens to the
+        # `unsplittable_tokens` arg of the superclass' constructor. Hence, we
+        # delete it from the config here.
+        del config["unsplittable_tokens"]
+
+        config.update(
+            {
+                "special_tokens": self.special_tokens,
+                "language_tokens": self.language_tokens,
+            }
+        )
+        return config

keras_hub/src/models/xlm_roberta/__init__.py

@@ -0,0 +1,26 @@
+# 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.
+
+from keras_hub.src.models.xlm_roberta.xlm_roberta_backbone import (
+    XLMRobertaBackbone,
+)
+from keras_hub.src.models.xlm_roberta.xlm_roberta_presets import (
+    backbone_presets,
+)
+from keras_hub.src.models.xlm_roberta.xlm_roberta_tokenizer import (
+    XLMRobertaTokenizer,
+)
+from keras_hub.src.utils.preset_utils import register_presets
+
+register_presets(backbone_presets, (XLMRobertaBackbone, XLMRobertaTokenizer))