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

keras_hub/src/models/bloom/bloom_preprocessor.py

@@ -1,178 +0,0 @@
-# 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 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.bloom.bloom_backbone import BloomBackbone
-from keras_hub.src.models.bloom.bloom_tokenizer import BloomTokenizer
-from keras_hub.src.models.preprocessor import Preprocessor
-from keras_hub.src.utils.tensor_utils import preprocessing_function
-
-
-@keras_hub_export("keras_hub.models.BloomPreprocessor")
-class BloomPreprocessor(Preprocessor):
-    """BLOOM preprocessing layer which tokenizes and packs inputs.
-
-    This preprocessing layer will do 2 things:
-
-    - Tokenize the inputs using the `tokenizer`.
-    - Construct a dictionary with keys `"token_ids"`, `"padding_mask"`, that can
-        be passed directly to a `keras_hub.models.BloomBackbone`.
-
-    This layer can be used directly with `tf.data.Dataset.map` to preprocess
-    string data in the `(x, y, sample_weight)` format used by
-    `keras.Model.fit`.
-
-    The call method of this layer accepts three arguments, `x`, `y`, and
-    `sample_weight`. `x` can be a python string or tensor representing a single
-    segment, a list of python strings representing a batch of single segments,
-    or a list of tensors representing multiple segments to be packed together.
-    `y` and `sample_weight` are both optional, can have any format, and will be
-    passed through unaltered.
-
-    Args:
-        tokenizer: A `keras_hub.models.BloomTokenizer` instance.
-        sequence_length: The length of the packed inputs.
-        add_start_token: If `True`, the preprocessor will prepend the tokenizer
-            start token to each input sequence.
-        add_end_token: If `True`, the preprocessor will append the tokenizer
-            end token to each input sequence.
-
-    Call arguments:
-        x: A string, `tf.Tensor` or list of python strings.
-        y: Any label data. Will be passed through unaltered.
-        sample_weight: Any label weight data. Will be passed through unaltered.
-        sequence_length: Pass to override the configured `sequence_length` of
-            the layer.
-
-    Examples:
-
-    Directly calling the layer on data.
-    ```python
-    preprocessor = keras_hub.models.BloomPreprocessor.from_preset(
-        "bloom_560m_multi"
-    )
-    # Tokenize and pack a single sentence.
-    preprocessor("The quick brown fox jumped.")
-
-    # Tokenize a batch of single sentences.
-    preprocessor(["The quick brown fox jumped.", "Call me Ishmael."])
-
-    # Custom vocabulary.
-    features = ["a quick fox.", "a fox quick."]
-    vocab = {"<pad>": 0, "<s>":1, "</s>":2, "a": 3, "Ġquick": 4, "Ġfox": 5}
-    merges = ["Ġ q", "u i", "c k", "ui ck", "Ġq uick"]
-    merges += ["Ġ f", "o x", "Ġf ox"]
-    tokenizer = keras_hub.models.BloomTokenizer(
-        vocabulary=vocab,
-        merges=merges,
-    )
-    preprocessor = keras_hub.models.BloomPreprocessor(tokenizer=tokenizer)
-    preprocessor("The quick brown fox jumped.")
-    ```
-
-    Mapping with `tf.data.Dataset`.
-    ```python
-    preprocessor = keras_hub.models.BloomPreprocessor.from_preset(
-        "bloom_560m_multi"
-    )
-
-    text = tf.constant(["The quick brown fox jumped.", "Call me Ishmael."])
-    label = tf.constant([1, 1])
-
-    # Map labeled single sentences.
-    ds = tf.data.Dataset.from_tensor_slices((text, label))
-    ds = ds.map(preprocessor, num_parallel_calls=tf.data.AUTOTUNE)
-
-    # Map unlabeled single sentences.
-    ds = tf.data.Dataset.from_tensor_slices(text)
-    ds = ds.map(preprocessor, num_parallel_calls=tf.data.AUTOTUNE)
-    ```
-    """
-
-    backbone_cls = BloomBackbone
-    tokenizer_cls = BloomTokenizer
-
-    def __init__(
-        self,
-        tokenizer,
-        sequence_length=1024,
-        add_start_token=True,
-        add_end_token=True,
-        **kwargs,
-    ):
-        super().__init__(**kwargs)
-        self.tokenizer = tokenizer
-        self.packer = None
-        self.sequence_length = sequence_length
-        self.add_start_token = add_start_token
-        self.add_end_token = add_end_token
-
-    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.
-        self.packer = StartEndPacker(
-            start_value=self.tokenizer.start_token_id,
-            end_value=self.tokenizer.end_token_id,
-            pad_value=self.tokenizer.pad_token_id,
-            sequence_length=self.sequence_length,
-            return_padding_mask=True,
-        )
-        self.built = True
-
-    @preprocessing_function
-    def call(
-        self,
-        x,
-        y=None,
-        sample_weight=None,
-        sequence_length=None,
-    ):
-        sequence_length = sequence_length or self.sequence_length
-        token_ids, padding_mask = self.packer(
-            self.tokenizer(x),
-            sequence_length=sequence_length,
-            add_start_value=self.add_start_token,
-            add_end_value=self.add_end_token,
-        )
-        x = {
-            "token_ids": token_ids,
-            "padding_mask": padding_mask,
-        }
-        return keras.utils.pack_x_y_sample_weight(x, y, sample_weight)
-
-    def get_config(self):
-        config = super().get_config()
-        config.update(
-            {
-                "sequence_length": self.sequence_length,
-                "add_start_token": self.add_start_token,
-                "add_end_token": self.add_end_token,
-            }
-        )
-        return config
-
-    @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/electra/electra_preprocessor.py

@@ -1,155 +0,0 @@
-# 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 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.electra.electra_backbone import ElectraBackbone
-from keras_hub.src.models.electra.electra_tokenizer import ElectraTokenizer
-from keras_hub.src.models.preprocessor import Preprocessor
-from keras_hub.src.utils.tensor_utils import preprocessing_function
-
-
-@keras_hub_export("keras_hub.models.ElectraPreprocessor")
-class ElectraPreprocessor(Preprocessor):
-    """A ELECTRA preprocessing layer which tokenizes and packs inputs.
-
-    This preprocessing layer will do three things:
-
-     1. Tokenize any number of input segments using the `tokenizer`.
-     2. Pack the inputs together using a `keras_hub.layers.MultiSegmentPacker`.
-       with the appropriate `"[CLS]"`, `"[SEP]"` and `"[PAD]"` tokens.
-     3. Construct a dictionary of with keys `"token_ids"` and `"padding_mask"`,
-       that can be passed directly to a ELECTRA model.
-
-    This layer can be used directly with `tf.data.Dataset.map` to preprocess
-    string data in the `(x, y, sample_weight)` format used by
-    `keras.Model.fit`.
-
-    Args:
-        tokenizer: A `keras_hub.models.ElectraTokenizer` instance.
-        sequence_length: The length of the packed inputs.
-        truncate: string. The algorithm to truncate a list of batched segments
-            to fit within `sequence_length`. The value can be either
-            `round_robin` or `waterfall`:
-            - `"round_robin"`: Available space is assigned one token at a
-                time in a round-robin fashion to the inputs that still need
-                some, until the limit is reached.
-            - `"waterfall"`: The allocation of the budget is done using a
-                "waterfall" algorithm that allocates quota in a
-                left-to-right manner and fills up the buckets until we run
-                out of budget. It supports an arbitrary number of segments.
-
-    Call arguments:
-        x: A tensor of single string sequences, or a tuple of multiple
-            tensor sequences to be packed together. Inputs may be batched or
-            unbatched. For single sequences, raw python inputs will be converted
-            to tensors. For multiple sequences, pass tensors directly.
-        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.ElectraPreprocessor.from_preset(
-        "electra_base_discriminator_en"
-    )
-    preprocessor(["The quick brown fox jumped.", "Call me Ishmael."])
-
-    # Custom vocabulary.
-    vocab = ["[UNK]", "[CLS]", "[SEP]", "[PAD]", "[MASK]"]
-    vocab += ["The", "quick", "brown", "fox", "jumped", "."]
-    tokenizer = keras_hub.models.ElectraTokenizer(vocabulary=vocab)
-    preprocessor = keras_hub.models.ElectraPreprocessor(tokenizer)
-    preprocessor("The quick brown fox jumped.")
-    ```
-
-    Mapping with `tf.data.Dataset`.
-    ```python
-    preprocessor = keras_hub.models.ElectraPreprocessor.from_preset(
-        "electra_base_discriminator_en"
-    )
-
-    first = tf.constant(["The quick brown fox jumped.", "Call me Ishmael."])
-    second = tf.constant(["The fox tripped.", "Oh look, a whale."])
-    label = tf.constant([1, 1])
-    # Map labeled single sentences.
-    ds = tf.data.Dataset.from_tensor_slices((first, label))
-    ds = ds.map(preprocessor, num_parallel_calls=tf.data.AUTOTUNE)
-
-
-    # Map unlabeled single sentences.
-    ds = tf.data.Dataset.from_tensor_slices(first)
-    ds = ds.map(preprocessor, num_parallel_calls=tf.data.AUTOTUNE)
-
-    # Map labeled sentence pairs.
-    ds = tf.data.Dataset.from_tensor_slices(((first, second), label))
-    ds = ds.map(preprocessor, num_parallel_calls=tf.data.AUTOTUNE)
-    # Map unlabeled sentence pairs.
-    ds = tf.data.Dataset.from_tensor_slices((first, second))
-
-    # Watch out for tf.data's default unpacking of tuples here!
-    # Best to invoke the `preprocessor` directly in this case.
-    ds = ds.map(
-        lambda first, second: preprocessor(x=(first, second)),
-        num_parallel_calls=tf.data.AUTOTUNE,
-    )
-    ```
-    """
-
-    backbone_cls = ElectraBackbone
-    tokenizer_cls = ElectraTokenizer
-
-    def __init__(
-        self,
-        tokenizer,
-        sequence_length=512,
-        truncate="round_robin",
-        **kwargs,
-    ):
-        super().__init__(**kwargs)
-        self.tokenizer = tokenizer
-        self.packer = MultiSegmentPacker(
-            start_value=self.tokenizer.cls_token_id,
-            end_value=self.tokenizer.sep_token_id,
-            pad_value=self.tokenizer.pad_token_id,
-            truncate=truncate,
-            sequence_length=sequence_length,
-        )
-
-    def get_config(self):
-        config = super().get_config()
-        config.update(
-            {
-                "sequence_length": self.packer.sequence_length,
-                "truncate": self.packer.truncate,
-            }
-        )
-        return config
-
-    @preprocessing_function
-    def call(self, x, y=None, sample_weight=None):
-        x = x if isinstance(x, tuple) else (x,)
-        x = tuple(self.tokenizer(segment) for segment in x)
-        token_ids, segment_ids = self.packer(x)
-        x = {
-            "token_ids": token_ids,
-            "segment_ids": segment_ids,
-            "padding_mask": token_ids != self.tokenizer.pad_token_id,
-        }
-        return keras.utils.pack_x_y_sample_weight(x, y, sample_weight)

keras_hub/src/models/falcon/falcon_preprocessor.py

@@ -1,180 +0,0 @@
-# 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 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.falcon.falcon_backbone import FalconBackbone
-from keras_hub.src.models.falcon.falcon_tokenizer import FalconTokenizer
-from keras_hub.src.models.preprocessor import Preprocessor
-from keras_hub.src.utils.tensor_utils import preprocessing_function
-
-
-@keras_hub_export("keras_hub.models.FalconPreprocessor")
-class FalconPreprocessor(Preprocessor):
-    """Falcon preprocessing layer which tokenizes and packs inputs.
-
-    This preprocessing layer will do 2 things:
-
-    - Tokenize the inputs using the `tokenizer`.
-    - Construct a dictionary with keys `"token_ids"`, `"padding_mask"`, that can
-        be passed directly to a `keras_hub.models.FalconBackbone`.
-
-    This layer can be used directly with `tf.data.Dataset.map` to preprocess
-    string data in the `(x, y, sample_weight)` format used by
-    `keras.Model.fit`.
-
-    The call method of this layer accepts three arguments, `x`, `y`, and
-    `sample_weight`. `x` can be a python string or tensor representing a single
-    segment, a list of python strings representing a batch of single segments,
-    or a list of tensors representing multiple segments to be packed together.
-    `y` and `sample_weight` are both optional, can have any format, and will be
-    passed through unaltered.
-
-    `FalconPreprocessor` forces the input to have only one segment, as Falcon is
-    mainly used for generation tasks. For tasks having multi-segment inputs
-    like "glue/mnli", please use a model designed for classification purposes
-    such as BERT or RoBERTa.
-
-    Args:
-        tokenizer: A `keras_hub.models.FalconTokenizer` instance.
-        sequence_length: The length of the packed inputs.
-        add_start_token: If `True`, the preprocessor will prepend the tokenizer
-            start token to each input sequence.
-        add_end_token: If `True`, the preprocessor will append the tokenizer
-            end token to each input sequence.
-
-    Call arguments:
-        x: A string, `tf.Tensor` or list of python strings.
-        y: Any label data. Will be passed through unaltered.
-        sample_weight: Any label weight data. Will be passed through unaltered.
-        sequence_length: Pass to override the configured `sequence_length` of
-            the layer.
-
-    Examples:
-
-    Directly calling the layer on data.
-    ```python
-    preprocessor = keras_hub.models.FalconPreprocessor.from_preset("falcon_rw_1b")
-
-    # Tokenize and pack a single sentence.
-    preprocessor("The quick brown fox jumped.")
-
-    # Tokenize a batch of single sentences.
-    preprocessor(["The quick brown fox jumped.", "Call me Ishmael."])
-
-    # Custom vocabulary.
-    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"]
-    tokenizer = keras_hub.models.FalconTokenizer(
-        vocabulary=vocab,
-        merges=merges,
-    )
-    preprocessor = keras_hub.models.FalconPreprocessor(tokenizer=tokenizer)
-    preprocessor("The quick brown fox jumped.")
-    ```
-
-    Mapping with `tf.data.Dataset`.
-    ```python
-    preprocessor = keras_hub.models.FalconPreprocessor.from_preset("falcon_rw_1b")
-
-    text = tf.constant(["The quick brown fox jumped.", "Call me Ishmael."])
-    label = tf.constant([1, 1])
-
-    # Map labeled single sentences.
-    ds = tf.data.Dataset.from_tensor_slices((text, label))
-    ds = ds.map(preprocessor, num_parallel_calls=tf.data.AUTOTUNE)
-
-    # Map unlabeled single sentences.
-    ds = tf.data.Dataset.from_tensor_slices(text)
-    ds = ds.map(preprocessor, num_parallel_calls=tf.data.AUTOTUNE)
-    ```
-    """
-
-    backbone_cls = FalconBackbone
-    tokenizer_cls = FalconTokenizer
-
-    def __init__(
-        self,
-        tokenizer,
-        sequence_length=1024,
-        add_start_token=True,
-        add_end_token=True,
-        **kwargs,
-    ):
-        super().__init__(**kwargs)
-        self.tokenizer = tokenizer
-        self.packer = None
-        self.sequence_length = sequence_length
-        self.add_start_token = add_start_token
-        self.add_end_token = add_end_token
-
-    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.
-        self.packer = StartEndPacker(
-            start_value=self.tokenizer.start_token_id,
-            end_value=self.tokenizer.end_token_id,
-            pad_value=self.tokenizer.pad_token_id,
-            sequence_length=self.sequence_length,
-            return_padding_mask=True,
-        )
-        self.built = True
-
-    @preprocessing_function
-    def call(
-        self,
-        x,
-        y=None,
-        sample_weight=None,
-        sequence_length=None,
-    ):
-        sequence_length = sequence_length or self.sequence_length
-        token_ids, padding_mask = self.packer(
-            self.tokenizer(x),
-            sequence_length=sequence_length,
-            add_start_value=self.add_start_token,
-            add_end_value=self.add_end_token,
-        )
-        x = {
-            "token_ids": token_ids,
-            "padding_mask": padding_mask,
-        }
-        return keras.utils.pack_x_y_sample_weight(x, y, sample_weight)
-
-    def get_config(self):
-        config = super().get_config()
-        config.update(
-            {
-                "sequence_length": self.sequence_length,
-                "add_start_token": self.add_start_token,
-                "add_end_token": self.add_end_token,
-            }
-        )
-        return config
-
-    @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/gemma/gemma_preprocessor.py

@@ -1,184 +0,0 @@
-# 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 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.gemma.gemma_backbone import GemmaBackbone
-from keras_hub.src.models.gemma.gemma_tokenizer import GemmaTokenizer
-from keras_hub.src.models.preprocessor import Preprocessor
-from keras_hub.src.utils.tensor_utils import preprocessing_function
-
-
-@keras_hub_export("keras_hub.models.GemmaPreprocessor")
-class GemmaPreprocessor(Preprocessor):
-    """Gemma preprocessing layer which tokenizes and packs inputs.
-
-    This preprocessing layer will do 2 things:
-
-    - Tokenize the inputs using the `tokenizer`.
-    - Construct a dictionary with keys `"token_ids"`, `"padding_mask"`, that can
-        be passed directly to a `keras_hub.models.GemmaBackbone`.
-
-    This layer can be used directly with `tf.data.Dataset.map` to preprocess
-    string data in the `(x, y, sample_weight)` format used by
-    `keras.Model.fit`.
-
-    The call method of this layer accepts three arguments, `x`, `y`, and
-    `sample_weight`. `x` can be a python string or tensor representing a single
-    segment, a list of python strings representing a batch of single segments,
-    or a list of tensors representing multiple segments to be packed together.
-    `y` and `sample_weight` are both optional, can have any format, and will be
-    passed through unaltered.
-
-    `GemmaPreprocessor` expects the input to have only one segment, as Gemma is
-    mainly used for generation tasks. For tasks having multi-segment inputs
-    please combine inputs into a single string input before passing to the
-    preprocessor layer.
-
-    Args:
-        tokenizer: A `keras_hub.models.GemmaTokenizer` instance.
-        sequence_length: The length of the packed inputs.
-        add_start_token: If `True`, the preprocessor will prepend the tokenizer
-            start token to each input sequence.
-        add_end_token: If `True`, the preprocessor will append the tokenizer
-            end token to each input sequence.
-
-    Call arguments:
-        x: A string, `tf.Tensor` or list of python strings.
-        y: Any label data. Will be passed through unaltered.
-        sample_weight: Any label weight data. Will be passed through unaltered.
-        sequence_length: Pass to override the configured `sequence_length` of
-            the layer.
-
-    Examples:
-
-    Directly calling the layer on data.
-    ```python
-    preprocessor = keras_hub.models.GemmaPreprocessor.from_preset(
-        "gemma_2b_en"
-    )
-
-    # Tokenize and pack a single sentence.
-    preprocessor("The quick brown fox jumped.")
-
-    # Tokenize a batch of sentences.
-    preprocessor(["The quick brown fox jumped.", "Call me Ishmael."])
-
-    # Custom vocabulary.
-    bytes_io = io.BytesIO()
-    ds = tf.data.Dataset.from_tensor_slices(["The quick brown fox jumped."])
-    sentencepiece.SentencePieceTrainer.train(
-        sentence_iterator=ds.as_numpy_iterator(),
-        model_writer=bytes_io,
-        vocab_size=8,
-        model_type="WORD",
-        pad_id=0,
-        bos_id=1,
-        eos_id=2,
-        unk_id=3,
-        pad_piece="<pad>",
-        bos_piece="<bos>",
-        eos_piece="<eos>",
-        unk_piece="<unk>",
-    )
-    tokenizer = keras_hub.models.GemmaTokenizer(
-        proto=bytes_io.getvalue(),
-    )
-    preprocessor = keras_hub.models.GemmaPreprocessor(tokenizer=tokenizer)
-    preprocessor("The quick brown fox jumped.")
-    ```
-
-    Apply preprocessing to a `tf.data.Dataset`.
-    ```python
-    preprocessor = keras_hub.models.GemmaPreprocessor.from_preset(
-        "gemma_2b_en"
-    )
-
-    text = tf.constant(["The quick brown fox jumped.", "Call me Ishmael."])
-    label = tf.constant([1, 1])
-
-    # Map labeled single sentences.
-    ds = tf.data.Dataset.from_tensor_slices((text, label))
-    ds = ds.map(preprocessor, num_parallel_calls=tf.data.AUTOTUNE)
-
-    # Map unlabeled single sentences.
-    ds = tf.data.Dataset.from_tensor_slices(text)
-    ds = ds.map(preprocessor, num_parallel_calls=tf.data.AUTOTUNE)
-    ```
-    """
-
-    backbone_cls = GemmaBackbone
-    tokenizer_cls = GemmaTokenizer
-
-    def __init__(
-        self,
-        tokenizer,
-        sequence_length=1024,
-        add_start_token=True,
-        add_end_token=True,
-        **kwargs,
-    ):
-        super().__init__(**kwargs)
-
-        self.tokenizer = tokenizer
-        self.sequence_length = sequence_length
-        self.add_start_token = add_start_token
-        self.add_end_token = add_end_token
-
-    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.
-        self.packer = StartEndPacker(
-            start_value=self.tokenizer.start_token_id,
-            end_value=self.tokenizer.end_token_id,
-            pad_value=self.tokenizer.pad_token_id,
-            sequence_length=self.sequence_length,
-            return_padding_mask=True,
-        )
-        self.built = True
-
-    @preprocessing_function
-    def call(
-        self,
-        x,
-        y=None,
-        sample_weight=None,
-        sequence_length=None,
-    ):
-        sequence_length = sequence_length or self.sequence_length
-        token_ids, padding_mask = self.packer(
-            self.tokenizer(x),
-            sequence_length=sequence_length,
-            add_start_value=self.add_start_token,
-            add_end_value=self.add_end_token,
-        )
-        x = {
-            "token_ids": token_ids,
-            "padding_mask": padding_mask,
-        }
-        return keras.utils.pack_x_y_sample_weight(x, y, sample_weight)
-
-    def get_config(self):
-        config = super().get_config()
-        config.update(
-            {
-                "sequence_length": self.sequence_length,
-                "add_start_token": self.add_start_token,
-                "add_end_token": self.add_end_token,
-            }
-        )
-        return config