keras-hub 0.20.0.dev1__py3-none-any.whl → 0.21.0.dev1__py3-none-any.whl

keras_hub/src/models/qwen_moe/qwen_moe_tokenizer.py

@@ -0,0 +1,46 @@
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.models.qwen_moe.qwen_moe_backbone import QwenMoeBackbone
+from keras_hub.src.tokenizers.byte_pair_tokenizer import BytePairTokenizer
+
+
+@keras_hub_export(
+    "keras_hub.tokenizers.QwenMoeTokenizer",
+)
+class QwenMoeTokenizer(BytePairTokenizer):
+    """Tokenizer for Qwen Moe model.
+
+    This tokenizer implements byte-pair encoding (BPE) for Qwen models,
+    handling special tokens like BOS (beginning of sequence) and EOS (end of
+    sequence).
+
+    Args:
+        vocabulary: Dictionary mapping tokens to token IDs, or path to
+            vocabulary file.
+        merges: List of BPE merges, or path to merges file.
+        bos_token: Beginning of sequence token. Defaults to None.
+        eos_token: End of sequence token. Defaults to "<|endoftext|>".
+        misc_special_tokens: Set of additional special tokens. Defaults to
+            empty set.
+    """
+
+    backbone_cls = QwenMoeBackbone
+
+    def __init__(
+        self,
+        vocabulary=None,
+        merges=None,
+        **kwargs,
+    ):
+        # Add EOS token
+        eos_token = "<|endoftext|>"
+        self._add_special_token(eos_token, "end_token")
+
+        self.start_token_id = None
+        self.start_token = None
+        self.pad_token_id = 0
+
+        super().__init__(
+            vocabulary=vocabulary,
+            merges=merges,
+            **kwargs,
+        )

keras_hub/src/models/retinanet/retinanet_image_converter.py

@@ -6,16 +6,3 @@ from keras_hub.src.models.retinanet.retinanet_backbone import RetinaNetBackbone
 @keras_hub_export("keras_hub.layers.RetinaNetImageConverter")
 class RetinaNetImageConverter(ImageConverter):
     backbone_cls = RetinaNetBackbone
-
-    def __init__(
-        self,
-        *args,
-        **kwargs,
-    ):
-        # TODO: update presets and remove these old config options. They were
-        # never needed.
-        if "norm_mean" in kwargs:
-            kwargs["offset"] = [-x for x in kwargs.pop("norm_mean")]
-        if "norm_std" in kwargs:
-            kwargs["scale"] = [1.0 / x for x in kwargs.pop("norm_std")]
-        super().__init__(*args, **kwargs)

keras_hub/src/models/retinanet/retinanet_presets.py

@@ -11,7 +11,7 @@ backbone_presets = {
             "params": 34121239,
             "path": "retinanet",
         },
-        "kaggle_handle": "kaggle://keras/retinanet/keras/retinanet_resnet50_fpn_coco/3",
+        "kaggle_handle": "kaggle://keras/retinanet/keras/retinanet_resnet50_fpn_coco/4",
     },
     "retinanet_resnet50_fpn_v2_coco": {
         "metadata": {
@@ -22,6 +22,6 @@ backbone_presets = {
             "params": 31558592,
             "path": "retinanet",
         },
-        "kaggle_handle": "kaggle://keras/retinanet/keras/retinanet_resnet50_fpn_v2_coco/2",
+        "kaggle_handle": "kaggle://keras/retinanet/keras/retinanet_resnet50_fpn_v2_coco/3",
     },
 }

keras_hub/src/models/segformer/segformer_image_segmenter_preprocessor.py

@@ -1,5 +1,3 @@
-import keras
-
 from keras_hub.src.api_export import keras_hub_export
 from keras_hub.src.models.image_segmenter_preprocessor import (
     ImageSegmenterPreprocessor,
@@ -8,25 +6,9 @@ from keras_hub.src.models.segformer.segformer_backbone import SegFormerBackbone
 from keras_hub.src.models.segformer.segformer_image_converter import (
     SegFormerImageConverter,
 )
-from keras_hub.src.utils.tensor_utils import preprocessing_function
-
-IMAGENET_DEFAULT_MEAN = [0.485, 0.456, 0.406]
-IMAGENET_DEFAULT_STD = [0.229, 0.224, 0.225]
 
 
 @keras_hub_export("keras_hub.models.SegFormerImageSegmenterPreprocessor")
 class SegFormerImageSegmenterPreprocessor(ImageSegmenterPreprocessor):
     backbone_cls = SegFormerBackbone
     image_converter_cls = SegFormerImageConverter
-
-    @preprocessing_function
-    def call(self, x, y=None, sample_weight=None):
-        if self.image_converter:
-            x = self.image_converter(x)
-            if y is not None:
-                y = self.image_converter(y)
-
-        x = x / 255
-        x = (x - IMAGENET_DEFAULT_MEAN) / IMAGENET_DEFAULT_STD
-
-        return keras.utils.pack_x_y_sample_weight(x, y, sample_weight)

keras_hub/src/models/segformer/segformer_presets.py

@@ -10,7 +10,7 @@ presets = {
             "params": 3719027,
             "path": "segformer_b0",
         },
-        "kaggle_handle": "kaggle://keras/segformer/keras/segformer_b0_ade20k_512/2",
+        "kaggle_handle": "kaggle://keras/segformer/keras/segformer_b0_ade20k_512/3",
     },
     "segformer_b1_ade20k_512": {
         "metadata": {
@@ -21,7 +21,7 @@ presets = {
             "params": 13682643,
             "path": "segformer_b1",
         },
-        "kaggle_handle": "kaggle://keras/segformer/keras/segformer_b1_ade20k_512/2",
+        "kaggle_handle": "kaggle://keras/segformer/keras/segformer_b1_ade20k_512/5",
     },
     "segformer_b2_ade20k_512": {
         "metadata": {
@@ -32,7 +32,7 @@ presets = {
             "params": 24727507,
             "path": "segformer_b2",
         },
-        "kaggle_handle": "kaggle://keras/segformer/keras/segformer_b2_ade20k_512/2",
+        "kaggle_handle": "kaggle://keras/segformer/keras/segformer_b2_ade20k_512/3",
     },
     "segformer_b3_ade20k_512": {
         "metadata": {
@@ -43,7 +43,7 @@ presets = {
             "params": 44603347,
             "path": "segformer_b3",
         },
-        "kaggle_handle": "kaggle://keras/segformer/keras/segformer_b3_ade20k_512/2",
+        "kaggle_handle": "kaggle://keras/segformer/keras/segformer_b3_ade20k_512/3",
     },
     "segformer_b4_ade20k_512": {
         "metadata": {
@@ -54,7 +54,7 @@ presets = {
             "params": 61373907,
             "path": "segformer_b4",
         },
-        "kaggle_handle": "kaggle://keras/segformer/keras/segformer_b4_ade20k_512/2",
+        "kaggle_handle": "kaggle://keras/segformer/keras/segformer_b4_ade20k_512/3",
     },
     "segformer_b5_ade20k_640": {
         "metadata": {
@@ -65,7 +65,7 @@ presets = {
             "params": 81974227,
             "path": "segformer_b5",
         },
-        "kaggle_handle": "kaggle://keras/segformer/keras/segformer_b5_ade20k_640/2",
+        "kaggle_handle": "kaggle://keras/segformer/keras/segformer_b5_ade20k_640/3",
     },
     "segformer_b0_cityscapes_1024": {
         "metadata": {
@@ -76,7 +76,7 @@ presets = {
             "params": 3719027,
             "path": "segformer_b0",
         },
-        "kaggle_handle": "kaggle://keras/segformer/keras/segformer_b0_cityscapes_1024/2",
+        "kaggle_handle": "kaggle://keras/segformer/keras/segformer_b0_cityscapes_1024/3",
     },
     "segformer_b1_cityscapes_1024": {
         "metadata": {
@@ -87,7 +87,7 @@ presets = {
             "params": 13682643,
             "path": "segformer_b1",
         },
-        "kaggle_handle": "kaggle://keras/segformer/keras/segformer_b1_ade20k_512/2",
+        "kaggle_handle": "kaggle://keras/segformer/keras/segformer_b1_ade20k_512/1",
     },
     "segformer_b2_cityscapes_1024": {
         "metadata": {
@@ -98,7 +98,7 @@ presets = {
             "params": 24727507,
             "path": "segformer_b2",
         },
-        "kaggle_handle": "kaggle://keras/segformer/keras/segformer_b2_cityscapes_1024/2",
+        "kaggle_handle": "kaggle://keras/segformer/keras/segformer_b2_cityscapes_1024/3",
     },
     "segformer_b3_cityscapes_1024": {
         "metadata": {
@@ -109,7 +109,7 @@ presets = {
             "params": 44603347,
             "path": "segformer_b3",
         },
-        "kaggle_handle": "kaggle://keras/segformer/keras/segformer_b3_cityscapes_1024/2",
+        "kaggle_handle": "kaggle://keras/segformer/keras/segformer_b3_cityscapes_1024/3",
     },
     "segformer_b4_cityscapes_1024": {
         "metadata": {
@@ -120,7 +120,7 @@ presets = {
             "params": 61373907,
             "path": "segformer_b4",
         },
-        "kaggle_handle": "kaggle://keras/segformer/keras/segformer_b4_cityscapes_1024/2",
+        "kaggle_handle": "kaggle://keras/segformer/keras/segformer_b4_cityscapes_1024/3",
     },
     "segformer_b5_cityscapes_1024": {
         "metadata": {
@@ -131,6 +131,6 @@ presets = {
             "params": 81974227,
             "path": "segformer_b5",
         },
-        "kaggle_handle": "kaggle://keras/segformer/keras/segformer_b5_cityscapes_1024/2",
+        "kaggle_handle": "kaggle://keras/segformer/keras/segformer_b5_cityscapes_1024/3",
     },
 }

keras_hub/src/models/stable_diffusion_3/stable_diffusion_3_backbone.py

@@ -550,6 +550,12 @@ class StableDiffusion3Backbone(Backbone):
         guidance_scale=None,
     ):
         step = ops.convert_to_tensor(step)
+        if not keras.utils.is_keras_tensor(num_steps):
+            num_steps = ops.convert_to_tensor(num_steps)
+        if guidance_scale is not None and not keras.utils.is_keras_tensor(
+            guidance_scale
+        ):
+            guidance_scale = ops.convert_to_tensor(guidance_scale)
         next_step = ops.add(step, 1)
         sigma, timestep = self.scheduler(step, num_steps)
         next_sigma, _ = self.scheduler(next_step, num_steps)

keras_hub/src/models/task.py

@@ -236,14 +236,17 @@ class Task(PipelineModel):
             objects_to_skip=backbone_layer_ids,
         )
 
-    def save_to_preset(self, preset_dir):
+    def save_to_preset(self, preset_dir, max_shard_size=10):
         """Save task to a preset directory.
 
         Args:
             preset_dir: The path to the local model preset directory.
+            max_shard_size: `int` or `float`. Maximum size in GB for each
+                sharded file. If `None`, no sharding will be done. Defaults to
+                `10`.
         """
         saver = get_preset_saver(preset_dir)
-        saver.save_task(self)
+        saver.save_task(self, max_shard_size=max_shard_size)
 
     @property
     def layers(self):

keras_hub/src/models/xception/__init__.py

@@ -0,0 +1,5 @@
+from keras_hub.src.models.xception.xception_backbone import XceptionBackbone
+from keras_hub.src.models.xception.xception_presets import backbone_presets
+from keras_hub.src.utils.preset_utils import register_presets
+
+register_presets(backbone_presets, XceptionBackbone)

keras_hub/src/models/xception/xception_backbone.py

@@ -0,0 +1,188 @@
+import functools
+
+from keras import layers
+
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.models.backbone import Backbone
+from keras_hub.src.utils.keras_utils import standardize_data_format
+
+
+@keras_hub_export("keras_hub.models.XceptionBackbone")
+class XceptionBackbone(Backbone):
+    """Xception core network with hyperparameters.
+
+    This class implements a Xception backbone as described in
+    [Xception: Deep Learning with Depthwise Separable Convolutions](https://arxiv.org/abs/1610.02357).
+
+    Most users will want the pretrained presets available with this model. If
+    you are creating a custom backbone, this model provides customizability
+    through the `stackwise_conv_filters` and `stackwise_pooling` arguments. This
+    backbone assumes the same basic structure as the original Xception mode:
+    * Residuals and pre-activation everywhere but the first and last block.
+    * Conv layers for the first block only, separable conv layers elsewhere.
+
+    Args:
+        stackwise_conv_filters: list of list of ints. Each outermost list
+            entry represents a block, and each innermost list entry a conv
+            layer. The integer value specifies the number of filters for the
+            conv layer.
+        stackwise_pooling: list of bools. A list of booleans per block, where
+            each entry is true if the block should includes a max pooling layer
+            and false if it should not.
+        image_shape: tuple. The input shape without the batch size.
+            Defaults to `(None, None, 3)`.
+        data_format: `None` or str. If specified, either `"channels_last"` or
+            `"channels_first"`. If unspecified, the Keras default will be used.
+        dtype: `None` or str or `keras.mixed_precision.DTypePolicy`. The dtype
+            to use for the model's computations and weights.
+
+    Examples:
+    ```python
+    input_data = np.random.uniform(0, 1, size=(2, 224, 224, 3))
+
+    # Pretrained Xception backbone.
+    model = keras_hub.models.Backbone.from_preset("xception_41_imagenet")
+    model(input_data)
+
+    # Randomly initialized Xception backbone with a custom config.
+    model = keras_hub.models.XceptionBackbone(
+        stackwise_conv_filters=[[32, 64], [64, 128], [256, 256]],
+        stackwise_pooling=[True, True, False],
+    )
+    model(input_data)
+    ```
+    """
+
+    def __init__(
+        self,
+        stackwise_conv_filters,
+        stackwise_pooling,
+        image_shape=(None, None, 3),
+        data_format=None,
+        dtype=None,
+        **kwargs,
+    ):
+        if len(stackwise_conv_filters) != len(stackwise_pooling):
+            raise ValueError("All stackwise args should have the same length.")
+
+        data_format = standardize_data_format(data_format)
+        channel_axis = -1 if data_format == "channels_last" else 1
+        num_blocks = len(stackwise_conv_filters)
+
+        # Layer shorcuts with common args.
+        norm = functools.partial(
+            layers.BatchNormalization,
+            axis=channel_axis,
+            dtype=dtype,
+        )
+        act = functools.partial(
+            layers.Activation,
+            activation="relu",
+            dtype=dtype,
+        )
+        conv = functools.partial(
+            layers.Conv2D,
+            kernel_size=(3, 3),
+            use_bias=False,
+            data_format=data_format,
+            dtype=dtype,
+        )
+        sep_conv = functools.partial(
+            layers.SeparableConv2D,
+            kernel_size=(3, 3),
+            padding="same",
+            use_bias=False,
+            data_format=data_format,
+            dtype=dtype,
+        )
+        point_conv = functools.partial(
+            layers.Conv2D,
+            kernel_size=(1, 1),
+            strides=(2, 2),
+            padding="same",
+            use_bias=False,
+            data_format=data_format,
+            dtype=dtype,
+        )
+        pool = functools.partial(
+            layers.MaxPool2D,
+            pool_size=(3, 3),
+            strides=(2, 2),
+            padding="same",
+            data_format=data_format,
+            dtype=dtype,
+        )
+
+        # === Functional Model ===
+        image_input = layers.Input(shape=image_shape)
+        x = image_input  # Intermediate result.
+
+        # Iterate through the blocks.
+        for block_i in range(num_blocks):
+            first_block, last_block = block_i == 0, block_i == num_blocks - 1
+            block_filters = stackwise_conv_filters[block_i]
+            use_pooling = stackwise_pooling[block_i]
+
+            # Save the block input as a residual.
+            residual = x
+            for conv_i, filters in enumerate(block_filters):
+                # First block has post activation and strides on first conv.
+                if first_block:
+                    prefix = f"block{block_i + 1}_conv{conv_i + 1}"
+                    strides = (2, 2) if conv_i == 0 else (1, 1)
+                    x = conv(filters, strides=strides, name=prefix)(x)
+                    x = norm(name=f"{prefix}_bn")(x)
+                    x = act(name=f"{prefix}_act")(x)
+                # Last block has post activation.
+                elif last_block:
+                    prefix = f"block{block_i + 1}_sepconv{conv_i + 1}"
+                    x = sep_conv(filters, name=prefix)(x)
+                    x = norm(name=f"{prefix}_bn")(x)
+                    x = act(name=f"{prefix}_act")(x)
+                else:
+                    prefix = f"block{block_i + 1}_sepconv{conv_i + 1}"
+                    # The first conv in second block has no activation.
+                    if block_i != 1 or conv_i != 0:
+                        x = act(name=f"{prefix}_act")(x)
+                    x = sep_conv(filters, name=prefix)(x)
+                    x = norm(name=f"{prefix}_bn")(x)
+
+            # Optional block pooling.
+            if use_pooling:
+                x = pool(name=f"block{block_i + 1}_pool")(x)
+
+            # Sum residual, first and last block do not have a residual.
+            if not first_block and not last_block:
+                prefix = f"block{block_i + 1}_residual"
+                filters = x.shape[channel_axis]
+                # Match filters with a pointwise conv if needed.
+                if filters != residual.shape[channel_axis]:
+                    residual = point_conv(filters, name=f"{prefix}_conv")(
+                        residual
+                    )
+                    residual = norm(name=f"{prefix}_bn")(residual)
+                x = layers.Add(name=f"{prefix}_add", dtype=dtype)([x, residual])
+
+        super().__init__(
+            inputs=image_input,
+            outputs=x,
+            dtype=dtype,
+            **kwargs,
+        )
+
+        # === Config ===
+        self.stackwise_conv_filters = stackwise_conv_filters
+        self.stackwise_pooling = stackwise_pooling
+        self.image_shape = image_shape
+        self.data_format = data_format
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "stackwise_conv_filters": self.stackwise_conv_filters,
+                "stackwise_pooling": self.stackwise_pooling,
+                "image_shape": self.image_shape,
+            }
+        )
+        return config

keras_hub/src/models/xception/xception_image_classifier.py

@@ -0,0 +1,12 @@
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.models.image_classifier import ImageClassifier
+from keras_hub.src.models.xception.xception_backbone import XceptionBackbone
+from keras_hub.src.models.xception.xception_image_classifier_preprocessor import (  # noqa: E501
+    XceptionImageClassifierPreprocessor,
+)
+
+
+@keras_hub_export("keras_hub.models.XceptionImageClassifier")
+class XceptionImageClassifier(ImageClassifier):
+    backbone_cls = XceptionBackbone
+    preprocessor_cls = XceptionImageClassifierPreprocessor

keras_hub/src/models/xception/xception_image_classifier_preprocessor.py

@@ -0,0 +1,14 @@
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.models.image_classifier_preprocessor import (
+    ImageClassifierPreprocessor,
+)
+from keras_hub.src.models.xception.xception_backbone import XceptionBackbone
+from keras_hub.src.models.xception.xception_image_converter import (
+    XceptionImageConverter,
+)
+
+
+@keras_hub_export("keras_hub.models.XceptionImageClassifierPreprocessor")
+class XceptionImageClassifierPreprocessor(ImageClassifierPreprocessor):
+    backbone_cls = XceptionBackbone
+    image_converter_cls = XceptionImageConverter

keras_hub/src/models/xception/xception_image_converter.py

@@ -0,0 +1,8 @@
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.layers.preprocessing.image_converter import ImageConverter
+from keras_hub.src.models.xception.xception_backbone import XceptionBackbone
+
+
+@keras_hub_export("keras_hub.layers.XceptionImageConverter")
+class XceptionImageConverter(ImageConverter):
+    backbone_cls = XceptionBackbone

keras_hub/src/models/xception/xception_presets.py

@@ -0,0 +1,14 @@
+"""Xception preset configurations."""
+
+backbone_presets = {
+    "xception_41_imagenet": {
+        "metadata": {
+            "description": (
+                "41-layer Xception model pre-trained on ImageNet 1k."
+            ),
+            "params": 20861480,
+            "path": "xception",
+        },
+        "kaggle_handle": "kaggle://keras/xception/keras/xception_41_imagenet/2",
+    },
+}

keras_hub/src/tests/mocks/mock_gemma3_tokenizer.py

@@ -0,0 +1,155 @@
+import tensorflow as tf
+
+from keras_hub.src.tokenizers.tokenizer 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
+
+
+class MockGemma3Tokenizer(Tokenizer):
+    def __init__(
+        self,
+        proto=None,
+        sequence_length=None,
+        dtype="int32",
+        add_bos=False,
+        add_eos=False,
+        **kwargs,
+    ):
+        if not is_int_dtype(dtype) and not is_string_dtype(dtype):
+            raise ValueError(
+                "Output dtype must be an integer type or a string. "
+                f"Received: dtype={dtype}"
+            )
+
+        super().__init__(dtype=dtype, **kwargs)
+
+        self.vocabulary = [
+            "<pad>",
+            "<bos>",
+            "<eos>",
+            "<unk>",
+            "<start_of_image>",
+            "<end_of_image>",
+            "<start_of_turn>",
+            "<end_of_turn>",
+            "<img>",
+            "the",
+            "brown",
+            "earth",
+            "fox",
+            "is",
+            "quick",
+            "round",
+            "\n\n",
+        ]
+        self.string_to_id = tf.lookup.StaticHashTable(
+            tf.lookup.KeyValueTensorInitializer(
+                self.vocabulary, list(range(len(self.vocabulary)))
+            ),
+            default_value=3,
+        )
+        self.id_to_string = tf.lookup.StaticHashTable(
+            tf.lookup.KeyValueTensorInitializer(
+                list(range(len(self.vocabulary))), self.vocabulary
+            ),
+            default_value="<unk>",
+        )
+
+        # The usual tokens.
+        self._add_special_token("<bos>", "start_token")
+        self._add_special_token("<eos>", "end_token")
+        self._add_special_token("<pad>", "pad_token")
+
+        # Image placeholder token.
+        self._add_special_token("<img>", "image_placeholder")
+
+        #  Some tokens which are used in the preprocessor. We need to keep them
+        # here so that the preprocessor works with `tf.data`.
+        self._add_special_token("<start_of_image>", "start_of_image_token")
+        self._add_special_token("<end_of_image>", "end_of_image_token")
+
+        # self.special_token_ids = [
+        #     0, 1, 2, 4, 5, 8
+        # ]
+
+        self.sequence_length = sequence_length
+        self.add_bos = add_bos
+        self.add_eos = add_eos
+
+    def vocabulary_size(self):
+        return len(self.vocabulary)
+
+    def get_vocabulary(self):
+        return self.vocabulary
+
+    def id_to_token(self, id):
+        return self.vocabulary[id]
+
+    def token_to_id(self, token):
+        return self.vocabulary.index(token)
+
+    @preprocessing_function
+    def tokenize(self, inputs):
+        inputs = tf.convert_to_tensor(inputs)
+        unbatched = inputs.shape.rank == 0
+        if unbatched:
+            inputs = tf.expand_dims(inputs, 0)
+
+        inputs = tf.strings.regex_replace(
+            inputs, self.start_of_image_token, f" {self.start_of_image_token} "
+        )
+        inputs = tf.strings.regex_replace(
+            inputs, self.end_of_image_token, f" {self.end_of_image_token} "
+        )
+        inputs = tf.strings.regex_replace(
+            inputs, self.image_placeholder, f" {self.image_placeholder} "
+        )
+        inputs = tf.strings.regex_replace(inputs, "  ", " ")
+
+        sep_inputs = tf.strings.split(inputs, sep=" ")
+        tokens = self.string_to_id.lookup(sep_inputs)
+
+        if self.add_bos:
+            bos_tensor = tf.fill(
+                value=self.start_token_id,
+                dims=tokens.shape.as_list()[0:1] + [1],
+            )
+            tokens = tf.concat((bos_tensor, tokens), axis=-1)
+        if self.add_eos:
+            eos_tensor = tf.fill(
+                value=self.end_token_id, dims=tokens.shape.as_list()[0:1] + [1]
+            )
+            tokens = tf.concat((tokens, eos_tensor), axis=-1)
+
+        # Convert to a dense output if input was a scalar.
+        if unbatched:
+            tokens = tf.squeeze(tokens, 0)
+
+        return tokens
+
+    @preprocessing_function
+    def detokenize(self, inputs):
+        inputs, unbatched, rectangular = convert_to_ragged_batch(inputs)
+        # tf-text sentencepiece does not handle int64.
+        inputs = tf.cast(inputs, "int32")
+
+        outputs = self.id_to_string.lookup(inputs)
+        outputs = tf.strings.reduce_join(outputs, axis=-1, separator=" ")
+
+        for token in [
+            self.start_token,
+            self.end_token,
+            self.pad_token,
+        ]:
+            outputs = tf.strings.regex_replace(outputs, token, "")
+
+        outputs = tf.strings.strip(outputs)
+
+        if unbatched:
+            outputs = tf.squeeze(outputs, 0)
+        return outputs
+
+    def __call__(self, inputs):
+        return self.tokenize(inputs)