keras-hub 0.21.1.dev0__py3-none-any.whl → 0.22.0.dev0__py3-none-any.whl

keras_hub/src/models/esm/esm_classifier_preprocessor.py

@@ -0,0 +1,135 @@
+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.esm.esm_backbone import ESMBackbone
+from keras_hub.src.models.esm.esm_tokenizer import ESMTokenizer
+from keras_hub.src.models.text_classifier_preprocessor import (
+    TextClassifierPreprocessor,
+)
+from keras_hub.src.utils.tensor_utils import preprocessing_function
+
+
+@keras_hub_export("keras_hub.models.ESMProteinClassifierPreprocessor")
+class ESMProteinClassifierPreprocessor(TextClassifierPreprocessor):
+    """A ESM 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.StartEndPacker`.
+       with the appropriate start, end and pad tokens.
+    3. Construct a dictionary with the key `"token_ids"`, that can be passed
+       directly to an ESM 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.ESMTokenizer` 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.ESMProteinClassifierPreprocessor.from_preset
+    (
+        hf://facebook/esm2_t6_8M_UR50D
+    )
+
+    # 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."])
+
+    # Preprocess a batch of sentence pairs.
+    # When handling multiple sequences, always convert to tensors first!
+    first = tf.constant(["The quick brown fox jumped.", "Call me Ishmael."])
+    second = tf.constant(["The fox tripped.", "Oh look, a whale."])
+    preprocessor((first, second))
+
+    # Custom vocabulary.
+    vocab = ["[UNK]", "[CLS]", "[SEP]", "[PAD]", "[MASK]"]
+    vocab += ["The", "quick", "brown", "fox", "jumped", "."]
+    tokenizer = keras_hub.models.ESMTokenizer(vocabulary=vocab)
+    preprocessor =
+        keras_hub.models.ESMProteinClassifierPreprocessor(tokenizer)
+    preprocessor("The quick brown fox jumped.")
+    ```
+
+    Mapping with `tf.data.Dataset`.
+    ```python
+    preprocessor = keras_hub.models.ESMProteinClassifierPreprocessor.from_preset
+    (
+        hf://facebook/esm2_t6_8M_UR50D
+    )
+
+    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 = ESMBackbone
+    tokenizer_cls = ESMTokenizer
+
+    def build(self, input_shape):
+        super().build(input_shape)
+        # Defer masker 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,
+        )
+
+    @preprocessing_function
+    def call(self, x, y=None, sample_weight=None):
+        x = self.tokenizer(x)
+        token_ids = self.packer(x)
+        x = {
+            "token_ids": token_ids,
+        }
+        return keras.utils.pack_x_y_sample_weight(x, y, sample_weight)

keras_hub/src/models/esm/esm_encoder.py

@@ -0,0 +1,134 @@
+import keras
+from keras import activations
+from keras import initializers
+
+from keras_hub.src.models.esm.esm_attention import EsmSelfAttention
+
+
+class ESMEncoder(keras.layers.Layer):
+    """MultiHeadAttention by ESM
+
+    Referred to the implementation of HuggingFace.
+    reference:
+        https://github.com/huggingface/transformers/
+        blob/main/src/transformers/models/esm/modeling_esm.py
+    """
+
+    def __init__(
+        self,
+        heads,
+        head_size,
+        intermediate_size=None,
+        max_wavelength=10000,
+        dropout=0,
+        activation="gelu",
+        use_bias=False,
+        kernel_initializer="glorot_uniform",
+        layer_norm_eps=1e-12,
+        use_rotary=True,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.heads = heads
+        self.head_size = head_size
+        self.intermediate_size = intermediate_size
+        self.use_bias = use_bias
+        self.kernel_initializer = initializers.get(kernel_initializer)
+        self.max_wavelength = max_wavelength
+        self.dropout = dropout
+        self.activation = activations.get(activation)
+        self.layer_norm_eps = layer_norm_eps
+        self.use_rotary = use_rotary
+
+    def build(self, input_shape):
+        super().build(input_shape)
+        self.attention_layer = EsmSelfAttention(
+            heads=self.heads,
+            head_size=self.head_size,
+            use_bias=self.use_bias,
+            max_wavelength=self.max_wavelength,
+            kernel_initializer=self.kernel_initializer,
+            dtype=self.dtype_policy,
+            use_rotary=self.use_rotary,
+            name="attention_layer",
+        )
+        self.attention_layer.build(input_shape)
+
+        self.dropout_layer = keras.layers.Dropout(
+            rate=self.dropout,
+            dtype=self.dtype_policy,
+            name="self_attention_dropout",
+        )
+        self.dropout_layer.build([])
+
+        # Feedforward layers.
+        self.feedforward_intermediate_dense = keras.layers.Dense(
+            self.intermediate_size,
+            kernel_initializer=self.kernel_initializer,
+            use_bias=self.use_bias,
+            dtype=self.dtype_policy,
+            activation=self.activation,
+            name="feedforward_intermediate_dense",
+        )
+        self.feedforward_intermediate_dense.build(input_shape)
+
+        self.feedforward_output_dense = keras.layers.Dense(
+            input_shape[-1],
+            kernel_initializer=self.kernel_initializer,
+            use_bias=self.use_bias,
+            dtype=self.dtype_policy,
+            name="feedforward_output_dense",
+        )
+
+        self.feedforward_output_dense.build(
+            [None, None, self.intermediate_size]
+        )
+
+        self.attention_norm = keras.layers.LayerNormalization(
+            epsilon=self.layer_norm_eps,
+            name="attention_norm",
+            dtype=self.dtype_policy,
+        )
+        self.attention_norm.build(input_shape)
+
+        self.feedforward_norm = keras.layers.LayerNormalization(
+            epsilon=self.layer_norm_eps,
+            name="ffn_norm",
+            dtype=self.dtype_policy,
+        )
+        self.feedforward_norm.build(input_shape)
+
+    def call(self, x, attention_mask=None):
+        attention_output = self.attention_layer(
+            self.attention_norm(self.dropout_layer(x)),
+            attention_mask=attention_mask,
+        )
+        residual = x + attention_output
+
+        x = self.feedforward_norm(self.dropout_layer(residual))
+        intermediate_output = self.feedforward_intermediate_dense(x)
+        feedforward_output = self.feedforward_output_dense(intermediate_output)
+        return residual + self.dropout_layer(feedforward_output)
+
+    def compute_output_shape(self, input_shape):
+        return input_shape
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "heads": self.heads,
+                "head_size": self.head_size,
+                "intermediate_size": self.intermediate_size,
+                "max_wavelength": self.max_wavelength,
+                "use_bias": self.use_bias,
+                "activation": activations.serialize(self.activation),
+                "dropout": self.dropout,
+                "layer_norm_eps": self.layer_norm_eps,
+                "use_rotary": self.use_rotary,
+                "kernel_initializer": initializers.serialize(
+                    self.kernel_initializer
+                ),
+            }
+        )
+        return config

keras_hub/src/models/esm/esm_masked_plm.py

@@ -0,0 +1,117 @@
+import keras
+
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.layers.modeling.masked_lm_head import MaskedLMHead
+from keras_hub.src.models.esm.esm_backbone import ESMBackbone
+from keras_hub.src.models.esm.esm_backbone import esm2_kernel_initializer
+from keras_hub.src.models.esm.esm_masked_plm_preprocessor import (
+    ESMMaskedPLMPreprocessor,
+)
+from keras_hub.src.models.masked_lm import MaskedLM
+
+
+@keras_hub_export(
+    ["keras_hub.models.ESM2MaskedPLM", "keras_hub.models.ESMMaskedPLM"]
+)
+class ESMMaskedPLM(MaskedLM):
+    """An end-to-end ESM2 model for the masked protein language modeling task.
+
+    This model will train ESM2 on a masked protein language modeling task.
+    The model will predict labels for a number of masked tokens in the
+    input data. For usage of this model with pre-trained weights, see the
+    `from_preset()` method.
+
+    This model can optionally be configured with a `preprocessor` layer, in
+    which case inputs can be raw string features during `fit()`, `predict()`,
+    and `evaluate()`. Inputs will be tokenized and dynamically masked during
+    training and evaluation. This is done by default when creating the model
+    with `from_preset()`.
+
+
+    Args:
+        backbone: A `keras_hub.models.ESM2Backbone` instance.
+        preprocessor: A `keras_hub.models.ESM2MaskedPLMPreprocessor` or
+            `None`. If `None`, this model will not apply preprocessing, and
+            inputs should be preprocessed before calling the model.
+
+    Examples:
+
+    Raw string data.
+    ```python
+    features = ["The quick brown fox jumped.", "I forgot my homework."]
+
+    # Pretrained protein language model.
+    masked_lm = keras_hub.models.ESM2MaskedPLM.from_preset(
+        "hf://facebook/esm2_t6_8M_UR50D",
+    )
+    masked_lm.fit(x=features, batch_size=2)
+
+    # Re-compile (e.g., with a new learning rate).
+    masked_lm.compile(
+        loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
+        optimizer=keras.optimizers.Adam(5e-5),
+        jit_compile=True,
+    )
+    # Access backbone programmatically (e.g., to change `trainable`).
+    masked_lm.backbone.trainable = False
+    # Fit again.
+    masked_lm.fit(x=features, batch_size=2)
+    ```
+
+    Preprocessed integer data.
+    ```python
+    # Create a preprocessed dataset where 0 is the mask token.
+    features = {
+        "token_ids": np.array([[1, 2, 0, 4, 0, 6, 7, 8]] * 2),
+        "mask_positions": np.array([[2, 4]] * 2)
+    }
+    # Labels are the original masked values.
+    labels = [[3, 5]] * 2
+
+    masked_lm = keras_hub.models.ESM2MaskedPLM.from_preset(
+        'hf://facebook/esm2_t6_8M_UR50D',
+        preprocessor=None,
+    )
+
+    masked_lm.fit(x=features, y=labels, batch_size=2)
+    ```
+    """
+
+    backbone_cls = ESMBackbone
+    preprocessor_cls = ESMMaskedPLMPreprocessor
+
+    def __init__(
+        self,
+        backbone,
+        preprocessor=None,
+        **kwargs,
+    ):
+        # === Layers ===
+        self.backbone = backbone
+        self.preprocessor = preprocessor
+        self.masked_lm_head = MaskedLMHead(
+            vocabulary_size=backbone.vocabulary_size,
+            intermediate_activation=backbone.activation,
+            kernel_initializer=esm2_kernel_initializer(),
+            dtype=backbone.dtype_policy,
+            layer_norm_epsilon=backbone.layer_norm_eps,
+            name="mlm_head",
+        )
+
+        # === Functional Model ===
+        inputs = {
+            **backbone.input,
+            "mask_positions": keras.Input(
+                shape=(None,), dtype="int32", name="mask_positions"
+            ),
+        }
+        backbone_outputs = backbone(backbone.input)
+
+        outputs = self.masked_lm_head(
+            backbone_outputs, inputs["mask_positions"]
+        )
+        super().__init__(
+            inputs=inputs,
+            outputs=outputs,
+            **kwargs,
+        )

keras_hub/src/models/esm/esm_masked_plm_preprocessor.py

@@ -0,0 +1,143 @@
+import keras
+
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.layers.preprocessing.masked_lm_mask_generator import (
+    MaskedLMMaskGenerator,
+)
+from keras_hub.src.layers.preprocessing.start_end_packer import StartEndPacker
+from keras_hub.src.models.esm.esm_backbone import ESMBackbone
+from keras_hub.src.models.esm.esm_tokenizer import ESMTokenizer
+from keras_hub.src.models.masked_lm_preprocessor import MaskedLMPreprocessor
+from keras_hub.src.utils.tensor_utils import preprocessing_function
+
+
+@keras_hub_export("keras_hub.models.ESMMaskedPLMPreprocessor")
+class ESMMaskedPLMPreprocessor(MaskedLMPreprocessor):
+    """ESM preprocessing for the masked language modeling task.
+
+    This preprocessing layer will prepare inputs for a masked language modeling
+    task. It is primarily intended for use with the
+    `keras_hub.models.ESMMaskedPLM` task model.
+    Preprocessing will occur in multiple steps.
+
+    1. Tokenize any number of input segments using the `tokenizer`.
+    2. Pack the inputs together with the appropriate `"[CLS]"`, `"[SEP]"` and
+      `"[PAD]"` tokens.
+    3. Randomly select non-special tokens to mask, controlled by
+      `mask_selection_rate`.
+    4. Construct a `(x, y, sample_weight)` tuple suitable for training with a
+      `keras_hub.models.ESMMaskedPLM` task model.
+
+    Args:
+        tokenizer: A `keras_hub.models.ESMTokenizer` instance.
+        sequence_length: int. 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.
+        mask_selection_rate: float. The probability an input token will be
+            dynamically masked.
+        mask_selection_length: int. The maximum number of masked tokens
+            in a given sample.
+        mask_token_rate: float. The probability the a selected token will be
+            replaced with the mask token.
+        random_token_rate: float. The probability the a selected token will be
+            replaced with a random token from the vocabulary. A selected token
+            will be left as is with probability
+            `1 - mask_token_rate - random_token_rate`.
+
+    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: Label data. Should always be `None` as the layer generates labels.
+        sample_weight: Label weights. Should always be `None` as the layer
+            generates label weights.
+
+    Examples:
+
+    Directly calling the layer on data.
+    ```python
+    preprocessor = keras_hub.models.ESMMaskedPLMPreprocessor.from_preset(
+        hf://facebook/esm2_t6_8M_UR50D
+    )
+
+    # Tokenize and mask a single sentence.
+    preprocessor("The quick brown fox jumped.")
+
+    # Tokenize and mask a batch of single sentences.
+    preprocessor(["The quick brown fox jumped.", "Call me Ishmael."])
+
+    # Tokenize and mask sentence pairs.
+    # In this case, always convert input to tensors before calling the layer.
+    first = tf.constant(["The quick brown fox jumped.", "Call me Ishmael."])
+    second = tf.constant(["The fox tripped.", "Oh look, a whale."])
+    preprocessor((first, second))
+    ```
+
+    Mapping with `tf.data.Dataset`.
+    ```python
+    preprocessor = keras_hub.models.ESMMaskedPLMPreprocessor.from_preset(
+        hf://facebook/esm2_t6_8M_UR50D
+    )
+
+    first = tf.constant(["The quick brown fox jumped.", "Call me Ishmael."])
+    second = tf.constant(["The fox tripped.", "Oh look, a whale."])
+
+    # Map single sentences.
+    ds = tf.data.Dataset.from_tensor_slices(first)
+    ds = ds.map(preprocessor, num_parallel_calls=tf.data.AUTOTUNE)
+
+    # Map 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 = ESMBackbone
+    tokenizer_cls = ESMTokenizer
+
+    def build(self, input_shape):
+        super().build(input_shape)
+        # Defer masker 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,
+        )
+        self.masker = MaskedLMMaskGenerator(
+            mask_selection_rate=self.mask_selection_rate,
+            mask_selection_length=self.mask_selection_length,
+            mask_token_rate=self.mask_token_rate,
+            random_token_rate=self.random_token_rate,
+            vocabulary_size=self.tokenizer.vocabulary_size(),
+            mask_token_id=self.tokenizer.mask_token_id,
+            unselectable_token_ids=self.tokenizer.special_token_ids,
+        )
+
+    @preprocessing_function
+    def call(self, x, y=None, sample_weight=None):
+        x = self.tokenizer(x)
+        token_ids = self.packer(x)
+        masker_outputs = self.masker(token_ids)
+        x = {
+            "token_ids": masker_outputs["token_ids"],
+            "mask_positions": masker_outputs["mask_positions"],
+        }
+        y = masker_outputs["mask_ids"]
+        sample_weight = masker_outputs["mask_weights"]
+        return keras.utils.pack_x_y_sample_weight(x, y, sample_weight)

keras_hub/src/models/esm/esm_presets.py

@@ -0,0 +1,53 @@
+"""ESM model preset configurations."""
+
+# Metadata for loading pretrained model weights.
+backbone_presets = {
+    "esm2_t6_8M": {
+        "metadata": {
+            "description": (
+                "6 transformer layers version of the ESM-2 protein language "
+                "model, trained on the UniRef50 clustered protein sequence "
+                "dataset."
+            ),
+            "params": 7_408_960,
+            "path": "esm",
+        },
+        "kaggle_handle": "kaggle://keras/esm-2/keras/esm2_t6_8M/1",
+    },
+    "esm2_t12_35M": {
+        "metadata": {
+            "description": (
+                "12 transformer layers version of the ESM-2 protein language "
+                "model, trained on the UniRef50 clustered protein sequence "
+                "dataset."
+            ),
+            "params": 33_269_280,
+            "path": "esm",
+        },
+        "kaggle_handle": "kaggle://keras/esm-2/keras/esm2_t12_35M/1",
+    },
+    "esm2_t30_150M": {
+        "metadata": {
+            "description": (
+                "30 transformer layers version of the ESM-2 protein language "
+                "model, trained on the UniRef50 clustered protein sequence "
+                "dataset."
+            ),
+            "params": 147_728_000,
+            "path": "esm",
+        },
+        "kaggle_handle": "kaggle://keras/esm-2/keras/esm2_t30_150M/1",
+    },
+    "esm2_t33_650M": {
+        "metadata": {
+            "description": (
+                "33 transformer layers version of the ESM-2 protein language "
+                "model, trained on the UniRef50 clustered protein sequence "
+                "dataset."
+            ),
+            "params": 649_400_320,
+            "path": "esm",
+        },
+        "kaggle_handle": "kaggle://keras/esm-2/keras/esm2_t33_650M/1",
+    },
+}

keras_hub/src/models/esm/esm_tokenizer.py

@@ -0,0 +1,82 @@
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.models.esm.esm_backbone import ESMBackbone
+from keras_hub.src.tokenizers.word_piece_tokenizer import WordPieceTokenizer
+
+
+@keras_hub_export(
+    [
+        "keras_hub.tokenizers.ESMTokenizer",
+        "keras_hub.models.ESMTokenizer",
+    ]
+)
+class ESMTokenizer(WordPieceTokenizer):
+    """A ESM tokenizer using WordPiece subword segmentation.
+
+    This tokenizer class will tokenize raw strings into integer sequences and
+    is based on `keras_hub.tokenizers.WordPieceTokenizer`. Unlike the
+    underlying tokenizer, it will check for special tokens needed by ESM
+    models and provides a `from_preset()` method to automatically download
+    a matching vocabulary for a ESM preset.
+
+    If input is a batch of strings (rank > 0), the layer will output a
+    `tf.RaggedTensor` where the last dimension of the output is ragged.
+
+    If input is a scalar string (rank == 0), the layer will output a dense
+    `tf.Tensor` with static shape `[None]`.
+
+    Args:
+        vocabulary: A list of strings or a string filename path. If
+            passing a list, each element of the list should be a single word
+            piece token string. If passing a filename, the file should be a
+            plain text file containing a single word piece token per line.
+        lowercase: If `True`, the input text will be first lowered before
+            tokenization.
+        special_tokens_in_strings: bool. A bool to indicate if the tokenizer
+            should expect special tokens in input strings that should be
+            tokenized and mapped correctly to their ids. Defaults to False.
+
+    Examples:
+    ```python
+    # Unbatched input.
+    tokenizer = keras_hub.models.ESMTokenizer.from_preset(
+        "hf://facebook/esm2_t6_8M_UR50D",
+    )
+    tokenizer("The quick brown fox jumped.")
+
+    # Batched input.
+    tokenizer(["The quick brown fox jumped.", "The fox slept."])
+
+    # Detokenization.
+    tokenizer.detokenize(tokenizer("The quick brown fox jumped."))
+
+    # Custom vocabulary.
+    vocab = ["[UNK]", "<cls>", "<eos>", "<pad>", "<mask>"]
+    vocab += ["The", "quick", "brown", "fox", "jumped", "."]
+    tokenizer = keras_hub.models.ESMTokenizer(vocabulary=vocab)
+    tokenizer("The quick brown fox jumped.")
+    ```
+    """
+
+    backbone_cls = ESMBackbone
+
+    def __init__(
+        self,
+        vocabulary=None,
+        lowercase=False,
+        oov_token="<unk>",
+        **kwargs,
+    ):
+        self._add_special_token("<cls>", "cls_token")
+        self._add_special_token("<eos>", "sep_token")
+        self._add_special_token("<pad>", "pad_token")
+        self._add_special_token("<mask>", "mask_token")
+        # Also add `tokenizer.start_token` and `tokenizer.end_token` for
+        # compatibility with other tokenizers.
+        self._add_special_token("<cls>", "start_token")
+        self._add_special_token("<eos>", "end_token")
+        super().__init__(
+            vocabulary=vocabulary,
+            lowercase=lowercase,
+            oov_token=oov_token,
+            **kwargs,
+        )

keras_hub/src/models/flux/flux_text_to_image_preprocessor.py

@@ -43,9 +43,13 @@ class FluxTextToImagePreprocessor(Preprocessor):
 
     def generate_preprocess(self, x):
         token_ids = {}
-        token_ids["clip_l"] = self.clip_l_preprocessor(x)["token_ids"]
+        token_ids["clip_l"] = self.clip_l_preprocessor(
+            {"prompts": x, "images": None}
+        )["token_ids"]
         if self.t5_preprocessor is not None:
-            token_ids["t5"] = self.t5_preprocessor(x)["token_ids"]
+            token_ids["t5"] = self.t5_preprocessor(
+                {"prompts": x, "images": None}
+            )["token_ids"]
         return token_ids
 
     def get_config(self):