keras-hub-nightly 0.15.0.dev20240823171555__py3-none-any.whl → 0.16.0.dev2024092017__py3-none-any.whl

keras_hub/src/models/xlm_roberta/{xlm_roberta_preprocessor.py → xlm_roberta_text_classifier_preprocessor.py}

@@ -19,17 +19,25 @@ from keras_hub.src.api_export import keras_hub_export
 from keras_hub.src.layers.preprocessing.multi_segment_packer import (
     MultiSegmentPacker,
 )
-from keras_hub.src.models.preprocessor import Preprocessor
+from keras_hub.src.models.text_classifier_preprocessor import (
+    TextClassifierPreprocessor,
+)
+from keras_hub.src.models.xlm_roberta.xlm_roberta_backbone import (
+    XLMRobertaBackbone,
+)
 from keras_hub.src.models.xlm_roberta.xlm_roberta_tokenizer import (
     XLMRobertaTokenizer,
 )
-from keras_hub.src.utils.keras_utils import (
-    convert_inputs_to_list_of_tensor_segments,
-)
+from keras_hub.src.utils.tensor_utils import preprocessing_function
 
 
-@keras_hub_export("keras_hub.models.XLMRobertaPreprocessor")
-class XLMRobertaPreprocessor(Preprocessor):
+@keras_hub_export(
+    [
+        "keras_hub.models.XLMRobertaTextClassifierPreprocessor",
+        "keras_hub.models.XLMRobertaPreprocessor",
+    ]
+)
+class XLMRobertaTextClassifierPreprocessor(TextClassifierPreprocessor):
     """An XLM-RoBERTa preprocessing layer which tokenizes and packs inputs.
 
     This preprocessing layer will do three things:
@@ -73,7 +81,7 @@ class XLMRobertaPreprocessor(Preprocessor):
 
     Directly calling the layer on data.
     ```python
-    preprocessor = keras_hub.models.XLMRobertaPreprocessor.from_preset(
+    preprocessor = keras_hub.models.TextClassifierPreprocessor.from_preset(
         "xlm_roberta_base_multi"
     )
 
@@ -107,13 +115,15 @@ class XLMRobertaPreprocessor(Preprocessor):
     )
     proto = train_sentencepiece(ds, vocab_size=10)
     tokenizer = keras_hub.models.XLMRobertaTokenizer(proto=proto)
-    preprocessor = keras_hub.models.XLMRobertaPreprocessor(tokenizer)
+    preprocessor = keras_hub.models.XLMRobertaTextClassifierPreprocessor(
+        tokenizer
+    )
     preprocessor("The quick brown fox jumped.")
     ```
 
     Mapping with `tf.data.Dataset`.
     ```python
-    preprocessor = keras_hub.models.XLMRobertaPreprocessor.from_preset(
+    preprocessor = keras_hub.models.TextClassifierPreprocessor.from_preset(
         "xlm_roberta_base_multi"
     )
 
@@ -144,25 +154,11 @@ class XLMRobertaPreprocessor(Preprocessor):
     ```
     """
 
+    backbone_cls = XLMRobertaBackbone
     tokenizer_cls = XLMRobertaTokenizer
 
-    def __init__(
-        self,
-        tokenizer,
-        sequence_length=512,
-        truncate="round_robin",
-        **kwargs,
-    ):
-        super().__init__(**kwargs)
-
-        self.tokenizer = tokenizer
-        self.packer = None
-        self.truncate = truncate
-        self.sequence_length = sequence_length
-
     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.
+        # Roberta is doubles up the sep token, so we override build.
         self.packer = MultiSegmentPacker(
             start_value=self.tokenizer.start_token_id,
             end_value=self.tokenizer.end_token_id,
@@ -173,33 +169,10 @@ class XLMRobertaPreprocessor(Preprocessor):
         )
         self.built = True
 
-    def get_config(self):
-        config = super().get_config()
-        config.update(
-            {
-                "sequence_length": self.sequence_length,
-                "truncate": self.truncate,
-            }
-        )
-        return config
-
+    @preprocessing_function
     def call(self, x, y=None, sample_weight=None):
-        x = convert_inputs_to_list_of_tensor_segments(x)
-        x = [self.tokenizer(segment) for segment in x]
-        token_ids, _ = self.packer(x)
-        x = {
-            "token_ids": token_ids,
-            "padding_mask": token_ids != self.tokenizer.pad_token_id,
-        }
+        output = super().call(x, y=y, sample_weight=sample_weight)
+        x, y, sample_weight = keras.utils.unpack_x_y_sample_weight(output)
+        # Backbone has no segment ID input.
+        del x["segment_ids"]
         return keras.utils.pack_x_y_sample_weight(x, y, sample_weight)
-
-    @property
-    def sequence_length(self):
-        """The padded length of model input sequences."""
-        return self._sequence_length
-
-    @sequence_length.setter
-    def sequence_length(self, value):
-        self._sequence_length = value
-        if self.packer is not None:
-            self.packer.sequence_length = value

keras_hub/src/models/xlm_roberta/xlm_roberta_tokenizer.py

@@ -14,6 +14,9 @@
 
 
 from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.models.xlm_roberta.xlm_roberta_backbone import (
+    XLMRobertaBackbone,
+)
 from keras_hub.src.tokenizers.sentence_piece_tokenizer import (
     SentencePieceTokenizer,
 )
@@ -25,7 +28,12 @@ except ImportError:
     tf = None
 
 
-@keras_hub_export("keras_hub.models.XLMRobertaTokenizer")
+@keras_hub_export(
+    [
+        "keras_hub.tokenizers.XLMRobertaTokenizer",
+        "keras_hub.models.XLMRobertaTokenizer",
+    ]
+)
 class XLMRobertaTokenizer(SentencePieceTokenizer):
     """An XLM-RoBERTa tokenizer using SentencePiece subword segmentation.
 
@@ -89,17 +97,24 @@ class XLMRobertaTokenizer(SentencePieceTokenizer):
     ```
     """
 
+    backbone_cls = XLMRobertaBackbone
+
     def __init__(self, proto, **kwargs):
-        # List of special tokens.
-        self._vocabulary_prefix = ["<s>", "<pad>", "</s>", "<unk>"]
+        # Handle special tokens manually, as the tokenizer maps these tokens in
+        # a way that is not reflected in the vocabulary.
+        self.start_token, self.start_token_id = "<s>", 0
+        self.pad_token, self.pad_token_id = "<pad>", 1
+        self.end_token, self.end_token_id = "</s>", 2
+        self.unk_token, self.unk_token_id = "<unk>", 3
+        super().__init__(proto=proto, **kwargs)
 
-        # IDs of special tokens.
-        self.start_token_id = 0  # <s>
-        self.pad_token_id = 1  # <pad>
-        self.end_token_id = 2  # </s>
-        self.unk_token_id = 3  # <unk>
+    @property
+    def special_tokens(self):
+        return ["<s>", "<pad>", "</s>", "<unk>"]
 
-        super().__init__(proto=proto, **kwargs)
+    @property
+    def special_token_ids(self):
+        return [0, 1, 2, 3]
 
     def set_proto(self, proto):
         super().set_proto(proto)
@@ -162,7 +177,7 @@ class XLMRobertaTokenizer(SentencePieceTokenizer):
 
         # Correct `unk_token_id` (0 -> 3). Note that we do not correct
         # `start_token_id` and `end_token_id`; they are dealt with in
-        # `XLMRobertaPreprocessor`.
+        # `XLMRobertaTextClassifierPreprocessor`.
         tokens = tf.where(tf.equal(tokens, 0), self.unk_token_id - 1, tokens)
 
         # Shift the tokens IDs right by one.

keras_hub/src/tests/test_case.py

@@ -18,6 +18,7 @@ import pathlib
 import re
 
 import keras
+import numpy as np
 import tensorflow as tf
 from absl.testing import parameterized
 from keras import ops
@@ -465,6 +466,8 @@ class TestCase(tf.test.TestCase, parameterized.TestCase):
         init_kwargs,
         input_data,
         expected_output_shape,
+        expected_pyramid_output_keys=None,
+        expected_pyramid_image_sizes=None,
         variable_length_data=None,
         run_mixed_precision_check=True,
         run_quantization_check=True,
@@ -492,6 +495,26 @@ class TestCase(tf.test.TestCase, parameterized.TestCase):
             run_quantization_check=run_quantization_check,
         )
 
+        if expected_pyramid_output_keys:
+            backbone = cls(**init_kwargs)
+            model = keras.models.Model(
+                backbone.inputs, backbone.pyramid_outputs
+            )
+            output_data = model(input_data)
+
+            self.assertIsInstance(output_data, dict)
+            self.assertEqual(
+                list(output_data.keys()), list(backbone.pyramid_outputs.keys())
+            )
+            self.assertEqual(
+                list(output_data.keys()), expected_pyramid_output_keys
+            )
+            # check height and width of each level.
+            for i, (k, v) in enumerate(output_data.items()):
+                self.assertEqual(
+                    tuple(v.shape[1:3]), expected_pyramid_image_sizes[i]
+                )
+
         # Check data_format. We assume that `input_data` is in "channels_last"
         # format.
         if run_data_format_check and can_run_data_format_check:
@@ -501,6 +524,12 @@ class TestCase(tf.test.TestCase, parameterized.TestCase):
                 input_data = ops.transpose(input_data, axes=(2, 0, 1))
             elif len(input_data_shape) == 4:
                 input_data = ops.transpose(input_data, axes=(0, 3, 1, 2))
+            if len(expected_output_shape) == 3:
+                x = expected_output_shape
+                expected_output_shape = (x[0], x[2], x[1])
+            elif len(expected_output_shape) == 4:
+                x = expected_output_shape
+                expected_output_shape = (x[0], x[3], x[1], x[2])
             if "image_shape" in init_kwargs:
                 init_kwargs = init_kwargs.copy()
                 init_kwargs["image_shape"] = tuple(
@@ -557,6 +586,10 @@ class TestCase(tf.test.TestCase, parameterized.TestCase):
         task.preprocessor = None
         task.fit(ds.map(preprocessor))
         task.preprocessor = preprocessor
+        # Turn off default compilation, should error during `fit()`.
+        task = cls(**init_kwargs, compile=False)
+        with self.assertRaisesRegex(ValueError, "You must call `compile"):
+            task.fit(ds)
 
     def run_preset_test(
         self,
@@ -567,6 +600,7 @@ class TestCase(tf.test.TestCase, parameterized.TestCase):
         expected_output=None,
         expected_output_shape=None,
         expected_partial_output=None,
+        expected_labels=None,
     ):
         """Run instantiation and a forward pass for a preset."""
         with self.assertRaises(Exception):
@@ -604,5 +638,17 @@ class TestCase(tf.test.TestCase, parameterized.TestCase):
 
             tree.map_structure(compare, output, expected_partial_output)
 
+        if expected_labels is not None:
+            output = ops.argmax(output, axis=-1)
+            self.assertAllEqual(output, expected_labels)
+
     def get_test_data_dir(self):
         return str(pathlib.Path(__file__).parent / "test_data")
+
+    def load_test_image(self, target_size=None):
+        # From https://commons.wikimedia.org/wiki/File:California_quail.jpg
+        path = os.path.join(self.get_test_data_dir(), "test_image.jpg")
+        img = keras.utils.load_img(
+            path, target_size=target_size, keep_aspect_ratio=True
+        )
+        return np.array(img)

keras_hub/src/tokenizers/byte_pair_tokenizer.py

@@ -31,6 +31,7 @@ from keras_hub.src.tokenizers import tokenizer
 from keras_hub.src.utils.tensor_utils import convert_to_ragged_batch
 from keras_hub.src.utils.tensor_utils import is_int_dtype
 from keras_hub.src.utils.tensor_utils import is_string_dtype
+from keras_hub.src.utils.tensor_utils import preprocessing_function
 
 try:
     import tensorflow as tf
@@ -63,12 +64,17 @@ def create_alts_for_unsplittable_tokens(unsplittable_tokens):
     # Create alternates for all special tokens that will be not split during
     # tokenization.
     alts = []
-    prefix = "Ĵ"
-    # Trim out splitters.
-    replace_pattern = r"'|\s+|[^\p{L}\p{N}]+"
-    for token in unsplittable_tokens:
-        token = re.sub(replace_pattern, "", token)
-        alts.append(prefix + token)
+    for index in range(len(unsplittable_tokens)):
+        # Map unsplittable tokens to ĴA, ĴB, ĴC, etc. Which we assume will be
+        # a very uncommon string in any input data. We can't use a literal
+        # numeric counter here because we will split on all numbers. Ĵ is a
+        # random character we chose as it is likely to be unique.
+        prefix = "Ĵ"
+        digits = [int(d) for d in str(index)]
+        # Make numbers to uppercase characters so our token is still
+        # unsplittable.
+        suffix = "".join([chr(ord("A") + d) for d in digits])
+        alts.append(prefix + suffix)
     return alts
 
 
@@ -252,9 +258,9 @@ class BytePairTokenizer(tokenizer.Tokenizer):
     array([1, 2], dtype=int32)
     >>> seq1, seq2 = tokenizer(["butterfly", "butter"])
     >>> np.array(seq1)
-    array([1, 2], dtype=int32)
+    array([1, 2])
     >>> np.array(seq2)
-    array([1], dtype=int32)
+    array([1])
     >>> tokenizer = keras_hub.tokenizers.BytePairTokenizer(
     ...     vocab, merge, sequence_length=2)
     >>> seq1, seq2 = tokenizer(["butterfly", "butter"])
@@ -268,8 +274,7 @@ class BytePairTokenizer(tokenizer.Tokenizer):
     >>> merge = ["b u", "t t", "e r", "bu tt", "butt er", "f l", "fl y"]
     >>> tokenizer = keras_hub.tokenizers.BytePairTokenizer(vocab, merge)
     >>> tokenizer.detokenize([[1, 2]])
-    <tf.Tensor: shape=(1,), dtype=string, numpy=array([b'butterfly'],
-    dtype=object)>
+    ['butterfly']
     """
 
     def __init__(
@@ -291,6 +296,8 @@ class BytePairTokenizer(tokenizer.Tokenizer):
         super().__init__(dtype=dtype, **kwargs)
         self.sequence_length = sequence_length
         self.add_prefix_space = add_prefix_space
+        if unsplittable_tokens is None:
+            unsplittable_tokens = self.special_tokens
         self.unsplittable_tokens = unsplittable_tokens
         self.file_assets = [VOCAB_FILENAME, MERGES_FILENAME]
 
@@ -385,6 +392,7 @@ class BytePairTokenizer(tokenizer.Tokenizer):
             list(range(len(self.merges))),
             default=self.merge_ranks_lookup_default,
         )
+        self._update_special_token_ids()
 
     def get_vocabulary(self):
         """Get the tokenizer vocabulary as a list of strings tokens."""
@@ -526,17 +534,21 @@ class BytePairTokenizer(tokenizer.Tokenizer):
                 "layer."
             )
 
+    @preprocessing_function
     def tokenize(self, inputs):
         self._check_vocabulary()
-        if not isinstance(inputs, (tf.Tensor, tf.RaggedTensor)):
-            inputs = tf.convert_to_tensor(inputs)
-
         if self.add_prefix_space:
             inputs = tf.strings.join([" ", inputs])
 
-        scalar_input = inputs.shape.rank == 0
-        if scalar_input:
+        inputs = tf.convert_to_tensor(inputs)
+        unbatched = inputs.shape.rank == 0
+        if unbatched:
             inputs = tf.expand_dims(inputs, 0)
+        if inputs.shape.rank > 1:
+            raise ValueError(
+                "`tokenize()` inputs should be a string, list of strings, or "
+                f"string tensor with rank < 2. Received: {inputs}"
+            )
 
         raw_tokens = split_strings_for_bpe(inputs, self.unsplittable_tokens)
         token_row_splits = raw_tokens.row_splits
@@ -581,15 +593,16 @@ class BytePairTokenizer(tokenizer.Tokenizer):
             tokens = tokens.to_tensor(shape=output_shape)
 
         # Convert to a dense output if input in scalar
-        if scalar_input:
+        if unbatched:
             tokens = tf.squeeze(tokens, 0)
             tf.ensure_shape(tokens, shape=[self.sequence_length])
 
         return tokens
 
+    @preprocessing_function
     def detokenize(self, inputs):
         self._check_vocabulary()
-        inputs, unbatched, _ = convert_to_ragged_batch(inputs)
+        inputs, unbatched, rectangular = convert_to_ragged_batch(inputs)
         inputs = tf.cast(inputs, self.dtype)
         unicode_text = tf.strings.reduce_join(
             self.id_to_token_map.lookup(inputs), axis=-1

keras_hub/src/tokenizers/byte_tokenizer.py

@@ -26,6 +26,7 @@ from keras_hub.src.api_export import keras_hub_export
 from keras_hub.src.tokenizers import tokenizer
 from keras_hub.src.utils.tensor_utils import convert_to_ragged_batch
 from keras_hub.src.utils.tensor_utils import is_int_dtype
+from keras_hub.src.utils.tensor_utils import preprocessing_function
 
 try:
     import tensorflow_text as tf_text
@@ -95,9 +96,9 @@ class ByteTokenizer(tokenizer.Tokenizer):
     >>> tokenizer = keras_hub.tokenizers.ByteTokenizer()
     >>> seq1, seq2 = tokenizer(inputs)
     >>> np.array(seq1)
-    array([104, 101, 108, 108, 111], dtype=int32)
+    array([104, 101, 108, 108, 111])
     >>> np.array(seq2)
-    array([104, 105], dtype=int32)
+    array([104, 105])
 
     Dense outputs.
     >>> inputs = ["hello", "hi"]
@@ -145,18 +146,16 @@ class ByteTokenizer(tokenizer.Tokenizer):
     Detokenization.
     >>> inputs = [104, 101, 108, 108, 111]
     >>> tokenizer = keras_hub.tokenizers.ByteTokenizer()
-    >>> outputs = tokenizer.detokenize(inputs)
-    >>> np.array(outputs).astype("U")
-    array('hello', dtype='<U5')
+    >>> tokenizer.detokenize(inputs)
+    'hello'
 
     Detokenization with invalid bytes.
     >>> # The 255 below is invalid utf-8.
     >>> inputs = [104, 101, 255, 108, 108, 111]
     >>> tokenizer = keras_hub.tokenizers.ByteTokenizer(
     ...     errors="replace", replacement_char=88)
-    >>> outputs = tokenizer.detokenize(inputs)
-    >>> np.array(outputs).astype("U")
-    array('heXllo', dtype='<U6')
+    >>> tokenizer.detokenize(inputs)
+    'heXllo'
     """
 
     def __init__(
@@ -201,6 +200,7 @@ class ByteTokenizer(tokenizer.Tokenizer):
         self._char_lst = tf.constant(
             [i.tobytes() for i in np.arange(256, dtype=np.uint8)]
         )
+        self._update_special_token_ids()
 
     def vocabulary_size(self):
         """Get the integer size of the tokenizer vocabulary."""
@@ -212,12 +212,10 @@ class ByteTokenizer(tokenizer.Tokenizer):
             vocab[chr(i)] = i
         return vocab
 
+    @preprocessing_function
     def tokenize(self, inputs):
-        if not isinstance(inputs, (tf.Tensor, tf.RaggedTensor)):
-            inputs = tf.convert_to_tensor(inputs)
-
-        scalar_input = inputs.shape.rank == 0
-        if scalar_input:
+        unbatched = inputs.shape.rank == 0
+        if unbatched:
             inputs = tf.expand_dims(inputs, 0)
 
         # Optional: Lowercase the input.
@@ -241,12 +239,13 @@ class ByteTokenizer(tokenizer.Tokenizer):
             output_shape[-1] = self.sequence_length
             tokens = tokens.to_tensor(shape=output_shape)
 
-        if scalar_input:
+        if unbatched:
             tokens = tf.squeeze(tokens, 0)
         return tokens
 
+    @preprocessing_function
     def detokenize(self, inputs):
-        inputs, unbatched, _ = convert_to_ragged_batch(inputs)
+        inputs, unbatched, rectangular = convert_to_ragged_batch(inputs)
         # Remove trailing padding tokens, so that trailing "\x00" bytes don't
         # show up in the detokenized output.
         inputs = tf.ragged.boolean_mask(inputs, tf.not_equal(inputs, 0))

keras_hub/src/tokenizers/sentence_piece_tokenizer.py

@@ -31,6 +31,7 @@ from keras_hub.src.tokenizers import tokenizer
 from keras_hub.src.utils.tensor_utils import convert_to_ragged_batch
 from keras_hub.src.utils.tensor_utils import is_int_dtype
 from keras_hub.src.utils.tensor_utils import is_string_dtype
+from keras_hub.src.utils.tensor_utils import preprocessing_function
 from keras_hub.src.utils.tensor_utils import tensor_to_list
 
 try:
@@ -66,6 +67,9 @@ class SentencePieceTokenizer(tokenizer.Tokenizer):
             for more details on the format.
         sequence_length: If set, the output will be converted to a dense
             tensor and padded/trimmed so all outputs are of `sequence_length`.
+        add_bos: Add beginning of sentence token to the result.
+        add_eos: Add end of sentence token to the result. Token is always
+            truncated if output is longer than specified `sequence_length`.
 
     References:
         - [Kudo and Richardson, 2018](https://arxiv.org/abs/1808.06226)
@@ -115,6 +119,8 @@ class SentencePieceTokenizer(tokenizer.Tokenizer):
         proto=None,
         sequence_length=None,
         dtype="int32",
+        add_bos=False,
+        add_eos=False,
         **kwargs,
     ) -> None:
         if not is_int_dtype(dtype) and not is_string_dtype(dtype):
@@ -127,6 +133,8 @@ class SentencePieceTokenizer(tokenizer.Tokenizer):
 
         self.proto = None
         self.sequence_length = sequence_length
+        self.add_bos = add_bos
+        self.add_eos = add_eos
         self.set_proto(proto)
         self.file_assets = [VOCAB_FILENAME]
 
@@ -171,10 +179,13 @@ class SentencePieceTokenizer(tokenizer.Tokenizer):
         self._sentence_piece = tf_text.SentencepieceTokenizer(
             model=proto_bytes,
             out_type=self.compute_dtype,
+            add_bos=self.add_bos,
+            add_eos=self.add_eos,
         )
         # Keras cannot serialize a bytestring, so we base64 encode the model
         # byte array as a string for saving.
         self.proto = proto_bytes
+        self._update_special_token_ids()
 
     def vocabulary_size(self):
         """Get the integer size of the tokenizer vocabulary."""
@@ -211,6 +222,8 @@ class SentencePieceTokenizer(tokenizer.Tokenizer):
             {
                 "proto": None,  # Save vocabulary via an asset!
                 "sequence_length": self.sequence_length,
+                "add_bos": self.add_bos,
+                "add_eos": self.add_eos,
             }
         )
         return config
@@ -222,12 +235,12 @@ class SentencePieceTokenizer(tokenizer.Tokenizer):
                 "sure to pass a `proto` argument when creating the layer."
             )
 
+    @preprocessing_function
     def tokenize(self, inputs):
         self._check_vocabulary()
-        if not isinstance(inputs, (tf.Tensor, tf.RaggedTensor)):
-            inputs = tf.convert_to_tensor(inputs)
-        scalar_input = inputs.shape.rank == 0
-        if scalar_input:
+        inputs = tf.convert_to_tensor(inputs)
+        unbatched = inputs.shape.rank == 0
+        if unbatched:
             inputs = tf.expand_dims(inputs, 0)
 
         if self._sentence_piece is None:
@@ -245,15 +258,15 @@ class SentencePieceTokenizer(tokenizer.Tokenizer):
             tokens = tokens.to_tensor(shape=output_shape)
 
         # Convert to a dense output if input was a scalar.
-        if scalar_input:
+        if unbatched:
             tokens = tf.squeeze(tokens, 0)
             tf.ensure_shape(tokens, shape=[self.sequence_length])
-
         return tokens
 
+    @preprocessing_function
     def detokenize(self, inputs):
         self._check_vocabulary()
-        inputs, unbatched, _ = convert_to_ragged_batch(inputs)
+        inputs, unbatched, rectangular = convert_to_ragged_batch(inputs)
         # tf-text sentencepiece does not handle int64.
         inputs = tf.cast(inputs, "int32")
         outputs = self._sentence_piece.detokenize(inputs)