keras-hub-nightly 0.16.1.dev202410020340__py3-none-any.whl → 0.19.0.dev202501260345__py3-none-any.whl

keras_hub/src/tests/test_case.py

@@ -15,7 +15,6 @@ from keras_hub.src.layers.modeling.reversible_embedding import (
 )
 from keras_hub.src.models.retinanet.feature_pyramid import FeaturePyramid
 from keras_hub.src.tokenizers.tokenizer import Tokenizer
-from keras_hub.src.utils.keras_utils import has_quantization_support
 from keras_hub.src.utils.tensor_utils import is_float_dtype
 
 
@@ -313,6 +312,14 @@ class TestCase(tf.test.TestCase, parameterized.TestCase):
 
         for policy in ["mixed_float16", "mixed_bfloat16", "bfloat16"]:
             policy = keras.mixed_precision.Policy(policy)
+            # Ensure the correct `dtype` is set for sublayers or submodels in
+            # `init_kwargs`.
+            original_init_kwargs = init_kwargs.copy()
+            for k, v in init_kwargs.items():
+                if isinstance(v, keras.Layer):
+                    config = v.get_config()
+                    config["dtype"] = policy
+                    init_kwargs[k] = v.__class__.from_config(config)
             layer = cls(**{**init_kwargs, "dtype": policy})
             if isinstance(layer, keras.Model):
                 output_data = layer(input_data)
@@ -343,8 +350,15 @@ class TestCase(tf.test.TestCase, parameterized.TestCase):
                     continue
                 self.assertEqual(policy.compute_dtype, sublayer.compute_dtype)
                 self.assertEqual(policy.variable_dtype, sublayer.variable_dtype)
+            # Restore `init_kwargs`.
+            init_kwargs = original_init_kwargs
 
     def run_quantization_test(self, instance, cls, init_kwargs, input_data):
+        # TODO: revert the following if. This works around a torch
+        # quantization failure in `MultiHeadAttention` with Keras 3.7.
+        if keras.config.backend() == "torch":
+            return
+
         def _get_supported_layers(mode):
             supported_layers = [keras.layers.Dense, keras.layers.EinsumDense]
             if mode == "int8":
@@ -361,6 +375,14 @@ class TestCase(tf.test.TestCase, parameterized.TestCase):
                     policy_map[layer.path] = keras.dtype_policies.get(
                         f"{mode}_from_float32"
                     )
+            # Ensure the correct `dtype` is set for sublayers or submodels in
+            # `init_kwargs`.
+            original_init_kwargs = init_kwargs.copy()
+            for k, v in init_kwargs.items():
+                if isinstance(v, keras.Layer):
+                    config = v.get_config()
+                    config["dtype"] = policy_map
+                    init_kwargs[k] = v.__class__.from_config(config)
             # Instantiate the layer.
             model = cls(**{**init_kwargs, "dtype": policy_map})
             # Call layer eagerly.
@@ -382,12 +404,16 @@ class TestCase(tf.test.TestCase, parameterized.TestCase):
             # Check weights loading.
             weights = model.get_weights()
             revived_model.set_weights(weights)
+            # Restore `init_kwargs`.
+            init_kwargs = original_init_kwargs
 
     def run_model_saving_test(
         self,
         cls,
         init_kwargs,
         input_data,
+        atol=0.000001,
+        rtol=0.000001,
     ):
         """Save and load a model from disk and assert output is unchanged."""
         model = cls(**init_kwargs)
@@ -401,7 +427,7 @@ class TestCase(tf.test.TestCase, parameterized.TestCase):
 
         # Check that output matches.
         restored_output = restored_model(input_data)
-        self.assertAllClose(model_output, restored_output)
+        self.assertAllClose(model_output, restored_output, atol=atol, rtol=rtol)
 
     def run_backbone_test(
         self,
@@ -431,8 +457,8 @@ class TestCase(tf.test.TestCase, parameterized.TestCase):
 
             # Check variable length sequences.
             if variable_length_data is None:
-                # If no variable length data passed, assume the second axis of all
-                # inputs is our sequence axis and create it ourselves.
+                # If no variable length data passed, assume the second axis of
+                # all inputs is our sequence axis and create it ourselves.
                 variable_length_data = [
                     tree.map_structure(
                         lambda x: x[:, :seq_length, ...], input_data
@@ -453,14 +479,14 @@ class TestCase(tf.test.TestCase, parameterized.TestCase):
         # Check name maps to classname.
         name = re.sub("(.)([A-Z][a-z]+)", r"\1_\2", cls.__name__)
         name = re.sub("([a-z])([A-Z])", r"\1_\2", name).lower()
-        self.assertRegexpMatches(backbone.name, name)
+        self.assertRegex(backbone.name, name)
 
         # Check mixed precision.
         if run_mixed_precision_check:
             self.run_precision_test(cls, init_kwargs, input_data)
 
         # Check quantization.
-        if run_quantization_check and has_quantization_support():
+        if run_quantization_check:
             self.run_quantization_test(backbone, cls, init_kwargs, input_data)
 
     def run_vision_backbone_test(
@@ -567,6 +593,15 @@ class TestCase(tf.test.TestCase, parameterized.TestCase):
         ds = tf.data.Dataset.from_tensor_slices(train_data).batch(batch_size)
         x, y, sw = keras.utils.unpack_x_y_sample_weight(train_data)
 
+        # Test: the tree struct output by the
+        # preprocessor must match what model expects.
+        preprocessed_data = preprocessor(*train_data)[0]
+        tree.assert_same_structure(
+            preprocessed_data,
+            task._inputs_struct,
+            check_types=False,
+        )
+
         # Test predict.
         output = task.predict(x)
         if expected_output_shape is not None:

keras_hub/src/tokenizers/byte_pair_tokenizer.py

@@ -43,7 +43,11 @@ SPLIT_PATTERN_1 = (
 SPLIT_PATTERN_1 = SPLIT_PATTERN_1.replace(
     "{special_spaces}", SPECIAL_WHITESPACES
 )
-SPLIT_PATTERN_2 = rf"""[\s६{SPECIAL_WHITESPACES}]$"""
+
+# The pattern " \t\r\f\v" is the same as \s "all spaces" but without the \n.
+# Multiple \n\n\n in sequence must not be split for Llama3.
+# SPLIT_PATTERN_2 = rf"""[\s६{SPECIAL_WHITESPACES}]$"""
+SPLIT_PATTERN_2 = rf"""[ \t\r\f\v६{SPECIAL_WHITESPACES}]$"""
 
 
 def create_alts_for_unsplittable_tokens(unsplittable_tokens):
@@ -196,8 +200,8 @@ class BytePairTokenizer(tokenizer.Tokenizer):
     """Bype-pair encoding tokenizer layer.
 
     This BPE tokenizer provides the same functionality as the official GPT-2
-    tokenizer. Given the same `vocabulary` which maps tokens to ids, and `merges`
-    which describes BPE merge rules, it should provide the same output
+    tokenizer. Given the same `vocabulary` which maps tokens to ids, and
+    `merges` which describes BPE merge rules, it should provide the same output
     as OpenAI implementation (https://github.com/openai/gpt-2/blob/master/src/encoder.py).
     Different from OpenAI, this implementation is graph-compatible, so you can
     use it within a `tf.data` pipeline.

keras_hub/src/tokenizers/byte_tokenizer.py

@@ -1,13 +1,5 @@
 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.tokenizers import tokenizer
 from keras_hub.src.utils.tensor_utils import convert_to_ragged_batch
@@ -15,8 +7,10 @@ from keras_hub.src.utils.tensor_utils import is_int_dtype
 from keras_hub.src.utils.tensor_utils import preprocessing_function
 
 try:
+    import tensorflow as tf
     import tensorflow_text as tf_text
 except ImportError:
+    tf = None
     tf_text = None
 
 
@@ -156,8 +150,7 @@ class ByteTokenizer(tokenizer.Tokenizer):
     ):
         if not is_int_dtype(dtype):
             raise ValueError(
-                "Output dtype must be an integer type. "
-                f"Received: dtype={dtype}"
+                f"Output dtype must be an integer type. Received: dtype={dtype}"
             )
 
         # Check normalization_form.

keras_hub/src/tokenizers/sentence_piece_tokenizer.py

@@ -4,14 +4,6 @@ import os
 
 import keras
 
-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.tokenizers import tokenizer
 from keras_hub.src.utils.tensor_utils import convert_to_ragged_batch
@@ -21,11 +13,12 @@ from keras_hub.src.utils.tensor_utils import preprocessing_function
 from keras_hub.src.utils.tensor_utils import tensor_to_list
 
 try:
+    import tensorflow as tf
     import tensorflow_text as tf_text
 except ImportError:
+    tf = None
     tf_text = None
 
-
 VOCAB_FILENAME = "vocabulary.spm"
 
 

keras_hub/src/tokenizers/sentence_piece_tokenizer_trainer.py

@@ -1,17 +1,13 @@
 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."
-    )
+from keras_hub.src.utils.tensor_utils import assert_tf_libs_installed
 
 try:
     import sentencepiece as spm
+    import tensorflow as tf
 except ImportError:
     spm = None
+    tf = None
 
 from keras_hub.src.api_export import keras_hub_export
 
@@ -52,7 +48,8 @@ def compute_sentence_piece_proto(
 
     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)
+    >>> 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:
@@ -82,6 +79,7 @@ def compute_sentence_piece_proto(
     tf.Tensor([ 4  8 12  5  9 14  5  6 13  4  7 10 11  6 13],
     shape=(15,), dtype=int32)
     """
+    assert_tf_libs_installed("compute_sentence_piece_proto")
 
     if spm is None:
         raise ImportError(
@@ -92,7 +90,8 @@ def compute_sentence_piece_proto(
 
     if not isinstance(data, (list, tuple, tf.data.Dataset)):
         raise ValueError(
-            "The `data` argument must be either `tf.data.Dataset` or `tuple` or `list`. "
+            "The `data` argument must be either `tf.data.Dataset` or "
+            "`tuple` or `list`. "
             f"Received: type(data)={type(data)}."
         )
 
@@ -105,8 +104,7 @@ def compute_sentence_piece_proto(
     model_writer = (
         open(proto_output_file, "wb") if proto_output_file else io.BytesIO()
     )
-    is_dataset = isinstance(data, tf.data.Dataset)
-    if is_dataset:
+    if tf is not None and isinstance(data, tf.data.Dataset):
         spm.SentencePieceTrainer.train(
             sentence_iterator=data.as_numpy_iterator(),
             model_writer=model_writer,

keras_hub/src/tokenizers/tokenizer.py

@@ -10,7 +10,7 @@ from keras_hub.src.utils.preset_utils import builtin_presets
 from keras_hub.src.utils.preset_utils import find_subclass
 from keras_hub.src.utils.preset_utils import get_file
 from keras_hub.src.utils.preset_utils import get_preset_loader
-from keras_hub.src.utils.preset_utils import save_serialized_object
+from keras_hub.src.utils.preset_utils import get_preset_saver
 from keras_hub.src.utils.python_utils import classproperty
 from keras_hub.src.utils.tensor_utils import preprocessing_function
 
@@ -66,7 +66,7 @@ class Tokenizer(PreprocessingLayer):
     backbone_cls = None
 
     def __init__(self, *args, **kwargs):
-        self.config_name = kwargs.pop("config_name", TOKENIZER_CONFIG_FILE)
+        self.config_file = kwargs.pop("config_file", TOKENIZER_CONFIG_FILE)
         super().__init__(*args, **kwargs)
         self.file_assets = None
 
@@ -178,7 +178,7 @@ class Tokenizer(PreprocessingLayer):
         config = super().get_config()
         config.update(
             {
-                "config_name": self.config_name,
+                "config_file": self.config_file,
             }
         )
         return config
@@ -189,11 +189,8 @@ class Tokenizer(PreprocessingLayer):
         Args:
             preset_dir: The path to the local model preset directory.
         """
-        save_serialized_object(self, preset_dir, config_file=self.config_name)
-        subdir = self.config_name.split(".")[0]
-        asset_dir = os.path.join(preset_dir, ASSET_DIR, subdir)
-        os.makedirs(asset_dir, exist_ok=True)
-        self.save_assets(asset_dir)
+        saver = get_preset_saver(preset_dir)
+        saver.save_tokenizer(self)
 
     @preprocessing_function
     def call(self, inputs, *args, training=None, **kwargs):
@@ -202,11 +199,11 @@ class Tokenizer(PreprocessingLayer):
     def load_preset_assets(self, preset):
         asset_path = None
         for asset in self.file_assets:
-            subdir = self.config_name.split(".")[0]
+            subdir = self.config_file.split(".")[0]
             preset_path = os.path.join(ASSET_DIR, subdir, asset)
             asset_path = get_file(preset, preset_path)
-        tokenizer_config_name = os.path.dirname(asset_path)
-        self.load_assets(tokenizer_config_name)
+        tokenizer_config_file = os.path.dirname(asset_path)
+        self.load_assets(tokenizer_config_file)
 
     @classproperty
     def presets(cls):
@@ -217,7 +214,7 @@ class Tokenizer(PreprocessingLayer):
     def from_preset(
         cls,
         preset,
-        config_name=TOKENIZER_CONFIG_FILE,
+        config_file=TOKENIZER_CONFIG_FILE,
         **kwargs,
     ):
         """Instantiate a `keras_hub.models.Tokenizer` from a model preset.
@@ -263,4 +260,4 @@ class Tokenizer(PreprocessingLayer):
         backbone_cls = loader.check_backbone_class()
         if cls.backbone_cls != backbone_cls:
             cls = find_subclass(preset, cls, backbone_cls)
-        return loader.load_tokenizer(cls, config_name, **kwargs)
+        return loader.load_tokenizer(cls, config_file, **kwargs)

keras_hub/src/tokenizers/unicode_codepoint_tokenizer.py

@@ -203,8 +203,7 @@ class UnicodeCodepointTokenizer(tokenizer.Tokenizer):
     ) -> None:
         if not is_int_dtype(dtype):
             raise ValueError(
-                "Output dtype must be an integer type. "
-                f"Received: dtype={dtype}"
+                f"Output dtype must be an integer type. Received: dtype={dtype}"
             )
 
         # Check normalization_form.
@@ -226,8 +225,9 @@ class UnicodeCodepointTokenizer(tokenizer.Tokenizer):
         if normalization_form:
             if input_encoding != "UTF-8":
                 raise ValueError(
-                    """Normalization Forms are Only Supported for Input Encoding
-                     UTF-8"""
+                    "Normalization Forms are Only Supported for Input "
+                    "Encoding UTF-8"
+                    ""
                 )
 
         super().__init__(dtype=dtype, **kwargs)
@@ -259,8 +259,9 @@ class UnicodeCodepointTokenizer(tokenizer.Tokenizer):
         return config
 
     def vocabulary_size(self):
-        """Get the size of the tokenizer vocabulary. None implies no vocabulary
-        size was provided"""
+        """Get the size of the tokenizer vocabulary.
+
+        None implies no vocabulary size was provided"""
         return self._vocabulary_size
 
     def get_vocabulary(self):
@@ -334,6 +335,7 @@ class UnicodeCodepointTokenizer(tokenizer.Tokenizer):
         id = ord(token)
         if id >= self.vocabulary_size():
             raise ValueError(
-                f"Token {token} is not supported by `UnicodeCodepointTokenizer`."
+                f"Token {token} is not supported by "
+                "`UnicodeCodepointTokenizer`."
             )
         return id

keras_hub/src/tokenizers/word_piece_tokenizer_trainer.py

@@ -1,5 +1,6 @@
 from keras_hub.src.api_export import keras_hub_export
 from keras_hub.src.tokenizers.word_piece_tokenizer import pretokenize
+from keras_hub.src.utils.tensor_utils import assert_tf_libs_installed
 
 try:
     import tensorflow as tf
@@ -55,7 +56,8 @@ def compute_word_piece_vocabulary(
         suffix_indicator: str. The characters prepended to a
             WordPiece to indicate that it is a suffix to another subword.
             E.g. `"##ing"`. Defaults to `"##"`.
-        reserved_tokens: list of strings. A list of tokens that must be included in the vocabulary.
+        reserved_tokens: list of strings. A list of tokens that must be included
+            in the vocabulary.
 
     Returns:
         Returns a list of vocabulary terms.
@@ -67,7 +69,10 @@ def compute_word_piece_vocabulary(
     >>> vocab = compute_word_piece_vocabulary(inputs, 13)
     >>> vocab
     ['[PAD]', '[CLS]', '[SEP]', '[UNK]', '[MASK]', 'a', 'b', 'm', 'p', 'r', 's', 't', '##at']
-    >>> tokenizer = keras_hub.tokenizers.WordPieceTokenizer(vocabulary=vocab, oov_token="[UNK]")
+    >>> tokenizer = keras_hub.tokenizers.WordPieceTokenizer(
+    ...     vocabulary=vocab,
+    ...     oov_token="[UNK]",
+    ... )
     >>> outputs = inputs.map(tokenizer.tokenize)
     >>> for x in outputs:
     ...     print(x)
@@ -112,7 +117,9 @@ def compute_word_piece_vocabulary(
     tokenizer = keras_hub.tokenizers.WordPieceTokenizer(vocabulary=vocab)
     inputs.map(tokenizer.tokenize)
     ```
-    """
+    """  # noqa: E501
+    assert_tf_libs_installed("compute_word_piece_vocabulary")
+
     # Read data files.
     if not isinstance(data, (list, tf.data.Dataset)):
         raise ValueError(

keras_hub/src/utils/keras_utils.py

@@ -2,7 +2,6 @@ import sys
 
 import keras
 from absl import logging
-from packaging.version import parse
 
 try:
     import tensorflow as tf
@@ -36,23 +35,13 @@ def print_msg(message, line_break=True):
         logging.info(message)
 
 
+# Register twice for backwards compat.
 @keras.saving.register_keras_serializable(package="keras_hub")
+@keras.saving.register_keras_serializable(package="keras_nlp")
 def gelu_approximate(x):
     return keras.activations.gelu(x, approximate=True)
 
 
-def has_quantization_support():
-    return False if parse(keras.version()) < parse("3.4.0") else True
-
-
-def assert_quantization_support():
-    if not has_quantization_support():
-        raise ValueError(
-            "Quantization API requires Keras >= 3.4.0 to function "
-            f"correctly. Received: '{keras.version()}'"
-        )
-
-
 def standardize_data_format(data_format):
     if data_format is None:
         return keras.config.image_data_format()

keras_hub/src/utils/pipeline_model.py

@@ -232,7 +232,7 @@ class PipelineModel(keras.Model):
     ):
         data = self.preprocess_samples(x, y, sample_weight)
         x, y, sample_weight = keras.utils.unpack_x_y_sample_weight(data)
-        x = ops.convert_to_tensor(x)
+        x = tree.map_structure(ops.convert_to_tensor, x)
         if y is not None:
             y = ops.convert_to_tensor(y)
         if sample_weight is not None:
@@ -253,7 +253,7 @@ class PipelineModel(keras.Model):
     ):
         data = self.preprocess_samples(x, y, sample_weight)
         x, y, sample_weight = keras.utils.unpack_x_y_sample_weight(data)
-        x = ops.convert_to_tensor(x)
+        x = tree.map_structure(ops.convert_to_tensor, x)
         if y is not None:
             y = ops.convert_to_tensor(y)
         if sample_weight is not None:
@@ -272,7 +272,7 @@ class PipelineModel(keras.Model):
     ):
         data = self.preprocess_samples(x)
         x, _, _ = keras.utils.unpack_x_y_sample_weight(data)
-        x = ops.convert_to_tensor(x)
+        x = tree.map_structure(ops.convert_to_tensor, x)
         return super().predict_on_batch(
             x=x,
             **kwargs,