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

keras_hub/src/tokenizers/sentence_piece_tokenizer_trainer.py

@@ -0,0 +1,150 @@
+# 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 io
+
+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."
+    )
+
+try:
+    import sentencepiece as spm
+except ImportError:
+    spm = None
+
+from keras_hub.src.api_export import keras_hub_export
+
+
+@keras_hub_export("keras_hub.tokenizers.compute_sentence_piece_proto")
+def compute_sentence_piece_proto(
+    data,
+    vocabulary_size,
+    model_type="unigram",
+    proto_output_file=None,
+    lowercase=False,
+):
+    r"""A utility to train a SentencePiece vocabulary.
+
+    Trains a SentencePiece vocabulary from an input dataset or a list of
+    filenames.
+
+    If `data` is a list of filenames, the file format is required to be plain
+    text files, and the text will be read in line by line during training.
+
+    Args:
+        data: A `tf.data.Dataset`, or a list of filenames.
+        vocabulary_size: int. The maximum size of a vocabulary to be trained.
+        model_type: str. The model algorithm must be one of
+            `"unigram"`, `"bpe"`, `"word"` or `"char"`. Defaults to `"unigram"`.
+        proto_output_file: str. If provided it will be used
+            as model_file which is passed to model_writer.
+            If `None`, the model_file will be `io.BytesIO` object.
+            Defaults to `None`.
+        lowercase: bool. If True, the input text will be
+            lowercased before tokenization. Defaults to `False`.
+
+    Returns:
+        A `bytes` object with a serialized SentencePiece proto or
+        `None` if proto_output_file if provided.
+
+    Examples:
+
+    Basic Usage (from Dataset).
+    >>> inputs = tf.data.Dataset.from_tensor_slices(["Drifting Along"])
+    >>> proto = keras_hub.tokenizers.compute_sentence_piece_proto(inputs, vocabulary_size=15)
+    >>> tokenizer = keras_hub.tokenizers.SentencePieceTokenizer(proto=proto)
+    >>> outputs = inputs.map(tokenizer)
+    >>> for output in outputs:
+    ...     print(output)
+    tf.Tensor([ 4  8 12  5  9 14  5  6 13  4  7 10 11  6 13],
+    shape=(15,), dtype=int32)
+
+    Basic Usage (with files).
+    ``` python
+    with open("test.txt", "w+") as f: f.write("Drifting Along\n")
+    inputs = ["test.txt"]
+    proto = keras_hub.tokenizers.compute_sentence_piece_proto(
+         inputs, vocabulary_size=15, proto_output_file="model.spm")
+    tokenizer = keras_hub.tokenizers.SentencePieceTokenizer(proto="model.spm")
+    ds = tf.data.Dataset.from_tensor_slices(["the quick brown fox."])
+    ds = ds.map(tokenizer)
+    ```
+
+    Usage with lowercase
+    >>> inputs = tf.data.Dataset.from_tensor_slices(["Drifting Along"])
+    >>> proto = keras_hub.tokenizers.compute_sentence_piece_proto(
+    ...     inputs, vocabulary_size=15, lowercase=True)
+    >>> tokenizer = keras_hub.tokenizers.SentencePieceTokenizer(proto=proto)
+    >>> outputs = inputs.map(tokenizer)
+    >>> for output in outputs:
+    ...     print(output)
+    tf.Tensor([ 4  8 12  5  9 14  5  6 13  4  7 10 11  6 13],
+    shape=(15,), dtype=int32)
+    """
+
+    if spm is None:
+        raise ImportError(
+            f"{compute_sentence_piece_proto.__name__} requires the "
+            "`sentencepiece` package. Please install it with "
+            "`pip install sentencepiece`."
+        )
+
+    if not isinstance(data, (list, tuple, tf.data.Dataset)):
+        raise ValueError(
+            "The `data` argument must be either `tf.data.Dataset` or `tuple` or `list`. "
+            f"Received: type(data)={type(data)}."
+        )
+
+    if model_type not in ["unigram", "bpe", "word", "char"]:
+        raise ValueError(
+            "The `model_type` argument must be one of `unigram`, `bpe`, `word`"
+            f"or `char`. Received: model_type={model_type}."
+        )
+
+    model_writer = (
+        open(proto_output_file, "wb") if proto_output_file else io.BytesIO()
+    )
+    is_dataset = isinstance(data, tf.data.Dataset)
+    if is_dataset:
+        spm.SentencePieceTrainer.train(
+            sentence_iterator=data.as_numpy_iterator(),
+            model_writer=model_writer,
+            vocab_size=vocabulary_size,
+            model_type=model_type,
+            normalization_rule_name="nmt_nfkc_cf" if lowercase else "nmt_nfkc",
+            pad_id=0,
+            unk_id=1,
+            bos_id=2,
+            eos_id=3,
+        )
+    else:
+        spm.SentencePieceTrainer.train(
+            input=data,
+            model_writer=model_writer,
+            vocab_size=vocabulary_size,
+            model_type=model_type,
+            normalization_rule_name="nmt_nfkc_cf" if lowercase else "nmt_nfkc",
+            pad_id=0,
+            unk_id=1,
+            bos_id=2,
+            eos_id=3,
+        )
+    if proto_output_file:
+        model_writer.close()
+    else:
+        return model_writer.getvalue()

keras_hub/src/tokenizers/tokenizer.py

@@ -0,0 +1,235 @@
+# 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 os
+
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.layers.preprocessing.preprocessing_layer import (
+    PreprocessingLayer,
+)
+from keras_hub.src.utils.preset_utils import TOKENIZER_ASSET_DIR
+from keras_hub.src.utils.preset_utils import TOKENIZER_CONFIG_FILE
+from keras_hub.src.utils.preset_utils import check_config_class
+from keras_hub.src.utils.preset_utils import check_format
+from keras_hub.src.utils.preset_utils import get_file
+from keras_hub.src.utils.preset_utils import list_presets
+from keras_hub.src.utils.preset_utils import list_subclasses
+from keras_hub.src.utils.preset_utils import load_serialized_object
+from keras_hub.src.utils.preset_utils import save_serialized_object
+from keras_hub.src.utils.preset_utils import save_tokenizer_assets
+from keras_hub.src.utils.python_utils import classproperty
+from keras_hub.src.utils.transformers.convert import load_transformers_tokenizer
+
+
+@keras_hub_export(
+    [
+        "keras_hub.models.Tokenizer",
+        "keras_hub.tokenizers.Tokenizer",
+    ]
+)
+class Tokenizer(PreprocessingLayer):
+    """A base class for tokenizer layers.
+
+    Tokenizers in the KerasHub library should all subclass this layer.
+    The class provides two core methods `tokenize()` and `detokenize()` for
+    going from plain text to sequences and back. A tokenizer is a subclass of
+    `keras.layers.Layer` and can be combined into a `keras.Model`.
+
+    Subclassers should always implement the `tokenize()` method, which will also
+    be the default when calling the layer directly on inputs.
+
+    Subclassers can optionally implement the `detokenize()` method if the
+    tokenization is reversible. Otherwise, this can be skipped.
+
+    Subclassers should implement `get_vocabulary()`, `vocabulary_size()`,
+    `token_to_id()` and `id_to_token()` if applicable. For some simple
+    "vocab free" tokenizers, such as a whitespace splitter show below, these
+    methods do not apply and can be skipped.
+
+    Example:
+
+    ```python
+    class WhitespaceSplitterTokenizer(keras_hub.tokenizers.Tokenizer):
+        def tokenize(self, inputs):
+            return tf.strings.split(inputs)
+
+        def detokenize(self, inputs):
+            return tf.strings.reduce_join(inputs, separator=" ", axis=-1)
+
+    tokenizer = WhitespaceSplitterTokenizer()
+
+    # Tokenize some inputs.
+    tokenizer.tokenize("This is a test")
+
+    # Shorthard for `tokenize()`.
+    tokenizer("This is a test")
+
+    # Detokenize some outputs.
+    tokenizer.detokenize(["This", "is", "a", "test"])
+    ```
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.file_assets = None
+
+    def tokenize(self, inputs, *args, **kwargs):
+        """Transform input tensors of strings into output tokens.
+
+        Args:
+            inputs: Input tensor, or dict/list/tuple of input tensors.
+            *args: Additional positional arguments.
+            **kwargs: Additional keyword arguments.
+        """
+        raise NotImplementedError(
+            "No implementation of `tokenize()` was found for "
+            f"{self.__class__.__name__}. All tokenizers should implement "
+            "`tokenize()`."
+        )
+
+    def detokenize(self, inputs, *args, **kwargs):
+        """Transform tokens back into strings.
+
+        Args:
+            inputs: Input tensor, or dict/list/tuple of input tensors.
+            *args: Additional positional arguments.
+            **kwargs: Additional keyword arguments.
+        """
+        raise NotImplementedError(
+            "No implementation of `detokenize()` was found for "
+            f"{self.__class__.__name__}."
+        )
+
+    def get_vocabulary(self):
+        """Get the tokenizer vocabulary as a list of strings terms."""
+        raise NotImplementedError(
+            "No implementation of `get_vocabulary()` was found for "
+            f"{self.__class__.__name__}."
+        )
+
+    def vocabulary_size(self):
+        """Returns the total size of the token id space."""
+        raise NotImplementedError(
+            "No implementation of `vocabulary_size()` was found for "
+            f"{self.__class__.__name__}."
+        )
+
+    def id_to_token(self, id):
+        """Convert an integer id to a string token."""
+        raise NotImplementedError(
+            "No implementation of `id_to_token()` was found for "
+            f"{self.__class__.__name__}."
+        )
+
+    def token_to_id(self, token):
+        """Convert a string token to an integer id."""
+        raise NotImplementedError(
+            "No implementation of `token_to_id()` was found for "
+            f"{self.__class__.__name__}."
+        )
+
+    def save_to_preset(self, preset_dir):
+        """Save tokenizer to a preset directory.
+
+        Args:
+            preset_dir: The path to the local model preset directory.
+        """
+        save_serialized_object(
+            self,
+            preset_dir,
+            config_file=TOKENIZER_CONFIG_FILE,
+        )
+        save_tokenizer_assets(self, preset_dir)
+
+    def call(self, inputs, *args, training=None, **kwargs):
+        return self.tokenize(inputs, *args, **kwargs)
+
+    def load_preset_assets(self, preset):
+        asset_path = None
+        for asset in self.file_assets:
+            asset_path = get_file(
+                preset, os.path.join(TOKENIZER_ASSET_DIR, asset)
+            )
+        tokenizer_asset_dir = os.path.dirname(asset_path)
+        self.load_assets(tokenizer_asset_dir)
+
+    @classproperty
+    def presets(cls):
+        """List built-in presets for a `Task` subclass."""
+        presets = list_presets(cls)
+        for subclass in list_subclasses(cls):
+            presets.update(subclass.presets)
+        return presets
+
+    @classmethod
+    def from_preset(
+        cls,
+        preset,
+        **kwargs,
+    ):
+        """Instantiate a `keras_hub.models.Tokenizer` from a model preset.
+
+        A preset is a directory of configs, weights and other file assets used
+        to save and load a pre-trained model. The `preset` can be passed as a
+        one of:
+
+        1. a built in preset identifier like `'bert_base_en'`
+        2. a Kaggle Models handle like `'kaggle://user/bert/keras/bert_base_en'`
+        3. a Hugging Face handle like `'hf://user/bert_base_en'`
+        4. a path to a local preset directory like `'./bert_base_en'`
+
+        For any `Tokenizer` subclass, you can run `cls.presets.keys()` to list
+        all built-in presets available on the class.
+
+        This constructor can be called in one of two ways. Either from the base
+        class like `keras_hub.models.Tokenizer.from_preset()`, or from
+        a model class like `keras_hub.models.GemmaTokenizer.from_preset()`.
+        If calling from the base class, the subclass of the returning object
+        will be inferred from the config in the preset directory.
+
+        Args:
+            preset: string. A built in preset identifier, a Kaggle Models
+                handle, a Hugging Face handle, or a path to a local directory.
+            load_weights: bool. If `True`, the weights will be loaded into the
+                model architecture. If `False`, the weights will be randomly
+                initialized.
+
+        Examples:
+        ```python
+        # Load a preset tokenizer.
+        tokenizer = keras_hub.tokenizerTokenizer.from_preset("bert_base_en")
+
+        # Tokenize some input.
+        tokenizer("The quick brown fox tripped.")
+
+        # Detokenize some input.
+        tokenizer.detokenize([5, 6, 7, 8, 9])
+        ```
+        """
+        format = check_format(preset)
+        if format == "transformers":
+            return load_transformers_tokenizer(cls, preset)
+
+        preset_cls = check_config_class(
+            preset, config_file=TOKENIZER_CONFIG_FILE
+        )
+        if not issubclass(preset_cls, cls):
+            raise ValueError(
+                f"Preset has type `{preset_cls.__name__}` which is not a "
+                f"a subclass of calling class `{cls.__name__}`. Call "
+                f"`from_preset` directly on `{preset_cls.__name__}` instead."
+            )
+
+        tokenizer = load_serialized_object(preset, TOKENIZER_CONFIG_FILE)
+        tokenizer.load_preset_assets(preset)
+        return tokenizer