keras-hub-nightly 0.22.0.dev202508100425__py3-none-any.whl → 0.22.0.dev202508120417__py3-none-any.whl

keras_hub/models/__init__.py

@@ -199,6 +199,22 @@ from keras_hub.src.models.electra.electra_backbone import (
 from keras_hub.src.models.electra.electra_tokenizer import (
     ElectraTokenizer as ElectraTokenizer,
 )
+from keras_hub.src.models.esm.esm_backbone import ESMBackbone as ESM2Backbone
+from keras_hub.src.models.esm.esm_backbone import ESMBackbone as ESMBackbone
+from keras_hub.src.models.esm.esm_classifier import (
+    ESMProteinClassifier as ESMProteinClassifier,
+)
+from keras_hub.src.models.esm.esm_classifier_preprocessor import (
+    ESMProteinClassifierPreprocessor as ESMProteinClassifierPreprocessor,
+)
+from keras_hub.src.models.esm.esm_masked_plm import (
+    ESMMaskedPLM as ESM2MaskedPLM,
+)
+from keras_hub.src.models.esm.esm_masked_plm import ESMMaskedPLM as ESMMaskedPLM
+from keras_hub.src.models.esm.esm_masked_plm_preprocessor import (
+    ESMMaskedPLMPreprocessor as ESMMaskedPLMPreprocessor,
+)
+from keras_hub.src.models.esm.esm_tokenizer import ESMTokenizer as ESMTokenizer
 from keras_hub.src.models.f_net.f_net_backbone import (
     FNetBackbone as FNetBackbone,
 )

keras_hub/src/models/esm/esm_attention.py

@@ -0,0 +1,95 @@
+import keras
+from keras import ops
+from packaging import version
+
+from keras_hub.src.layers.modeling.rotary_embedding import RotaryEmbedding
+from keras_hub.src.models.roformer_v2.roformer_v2_attention import (
+    RoformerAttention,
+)
+
+
+class ESMRotaryEmbedding(RotaryEmbedding):
+    def _compute_cos_sin_embedding(self, x, position=1):
+        dim = x.shape[-1]
+        inv_freq = self.scaling_factor / (
+            self.max_wavelength ** (ops.arange(0, dim, 2, dtype=x.dtype) / dim)
+        )
+        t = ops.arange(x.shape[position], dtype=x.dtype)
+        freqs = ops.outer(t, inv_freq)
+        emb = ops.concatenate((freqs, freqs), axis=-1)
+
+        cos_emb = ops.cos(emb)[None, :, None, :]
+        sin_emb = ops.sin(emb)[None, :, None, :]
+        return cos_emb, sin_emb
+
+    def call(self, q, k, position=1):
+        cos_emb, sin_emb = self._compute_cos_sin_embedding(q, position)
+
+        return (
+            self.apply_rotary_pos_emb(q, cos_emb, sin_emb),
+            self.apply_rotary_pos_emb(k, cos_emb, sin_emb),
+        )
+
+    def rotate_half(self, x):
+        x1, x2 = ops.split(x, 2, -1)
+        return ops.concatenate((-x2, x1), axis=-1)
+
+    def apply_rotary_pos_emb(self, x, cos, sin):
+        cos = cos[:, : x.shape[1], :, :]
+        sin = sin[:, : x.shape[1], :, :]
+
+        return (x * cos) + (self.rotate_half(x) * sin)
+
+
+class EsmSelfAttention(RoformerAttention):
+    """MultiHeadAttention by ESM2
+
+    Referred to the implementation of HuggingFace.
+    In fact, this part of the calculation is exactly the same as RoFormer.
+    Only the calculation of the rotary part is different.
+    """
+
+    def __init__(self, use_rotary=True, **kwargs):
+        super().__init__(**kwargs)
+        self.use_rotary = use_rotary
+
+    def build(self, input_shape):
+        super().build(input_shape)
+        if self.use_rotary:
+            self.rotary_embedding_layer = ESMRotaryEmbedding(
+                max_wavelength=self.max_wavelength, dtype=self.dtype_policy
+            )
+            self.rotary_embedding_layer.build([])
+
+    def call(self, x, attention_mask=None):
+        qw = self.q_dense(x)
+        kw = self.k_dense(x)
+        vw = self.v_dense(x)
+
+        b, s = ops.shape(qw)[:2]
+        qw = ops.reshape(qw, (b, s, self.heads, self.head_size))
+        kw = ops.reshape(kw, (b, s, self.heads, self.head_size))
+        vw = ops.reshape(vw, (b, s, self.heads, self.head_size))
+
+        if self.use_rotary:
+            qw, kw = self.rotary_embedding_layer(qw, kw)
+        if version.parse(keras.__version__) < version.parse("3.6"):
+            raise ValueError("Please make sure your Keras version is >=3.6.")
+        flash_attention = keras.config.is_flash_attention_enabled()
+        attention_mask = ops.reshape(attention_mask, [b, 1, s, 1])
+        if keras.config.backend() == "torch":
+            attention_mask = ops.repeat(attention_mask, s, -1)
+            attention_mask = ops.transpose(attention_mask, [0, 1, 3, 2])
+        o = ops.dot_product_attention(
+            qw, kw, vw, mask=attention_mask, flash_attention=flash_attention
+        )
+        return self.o_dense(ops.reshape(o, [b, s, -1]))
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "use_rotary": self.use_rotary,
+            }
+        )
+        return config

keras_hub/src/models/esm/esm_backbone.py

@@ -0,0 +1,229 @@
+import keras
+from keras import activations
+
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.layers.modeling.position_embedding import PositionEmbedding
+from keras_hub.src.models.backbone import Backbone
+from keras_hub.src.models.esm.esm_encoder import ESMEncoder
+
+
+def esm2_kernel_initializer(stddev=0.02):
+    return keras.initializers.TruncatedNormal(stddev=stddev)
+
+
+@keras_hub_export(
+    ["keras_hub.models.ESM2Backbone", "keras_hub.models.ESMBackbone"]
+)
+class ESMBackbone(Backbone):
+    """A ESM2 and ESM encoder network.
+
+    This class implements a bi-directional Transformer-based encoder as
+    described in ["ESM"](https://github.com/facebookresearch/esm).
+
+    The default constructor gives a fully customizable, randomly initialized
+    ESM2 encoder with any number of layers, heads, and embed dim.To
+    load preset architectures and weights, use the `from_preset()` constructor.
+
+
+    Args:
+        vocabulary_size: int. The size of the token vocabulary.
+        num_layers: int. The number of transformer layers.
+        num_heads: int. The number of attention heads for each transformer.
+            The hidden size must be divisible by the number of attention heads.
+        hidden_dim: int. The size of the transformer encoding and pooler layers.
+        intermediate_dim: int. The output dimension of the first Dense layer in
+            a two-layer feedforward network for each transformer.
+        dropout: float. Dropout probability for the Transformer encoder.
+            Defaults to 0.1
+        use_pre_layer_norm:bool.If true, then layer norm will be used before
+            entering the transformer block.
+            Since it's pre-norm, the default is false.
+        max_sequence_length: int. The maximum sequence length that this encoder
+            can consume. If None, `max_sequence_length` uses the value from
+            sequence length. This determines the variable shape for positional
+            embeddings.
+        position_embedding_type: str. The position embedding type to use.
+            One of "absolute" and "rotary".
+            Use "absolute" for ESM1. Use "rotary" for ESM2. Defaults to "rotary"
+        max_wavelength : int. The maximum  angular wavelength of
+            the sine/cosine curves, for rotary embeddings.
+            Defaults to `10000`.
+        activation :string or keras.activations. The activation to
+            use for the transformer.
+            Defaults to `"gelu"`.
+        pad_token_id: int.padding token id. Normally 0,
+            but is set to 1 in the esm2 model.
+            Defaults to 0.
+        dtype: None or str or keras.mixed_precision.DTypePolicy. The dtype to
+            use for model computations and weights. Note that some computations,
+            such as softmax and layer normalization, will always be done at
+            float32 precision regardless of dtype.
+
+    Examples:
+    ```python
+    input_data = {
+        "token_ids": np.ones(shape=(1, 12), dtype="int32"),
+    }
+
+    # Pretrained ESM2 encoder.
+    model = keras_hub.models.ESM2Backbone.from_preset('hf://facebook/esm2_t6_8M_UR50D')
+    model(input_data)
+
+    # Randomly initialized ESM2 encoder with a custom config.
+    model = keras_hub.models.ESM2Backbone(
+        vocabulary_size=30552,
+        num_layers=4,
+        num_heads=4,
+        hidden_dim=256,
+        intermediate_dim=512,
+    )
+    model(input_data)
+    ```
+    """
+
+    def __init__(
+        self,
+        vocabulary_size,
+        num_layers,
+        num_heads,
+        hidden_dim,
+        intermediate_dim,
+        use_bias=True,
+        activation="gelu",
+        dropout=0.1,
+        dtype=None,
+        max_sequence_length=1024,
+        max_wavelength=10000,
+        layer_norm_eps=1e-12,
+        use_pre_layer_norm=False,
+        position_embedding_type="rotary",
+        pad_token_id=0,
+        **kwargs,
+    ):
+        if position_embedding_type not in (
+            "rotary",
+            "absolute",
+        ):
+            raise ValueError(
+                '`position_embedding_type` must be either `"rotary"`, or '
+                '`"absolute"`. Received '
+                f"position_embedding_type={position_embedding_type}."
+            )
+        head_size = hidden_dim // num_heads
+        # === Layers ===
+        self.token_embedding = keras.layers.Embedding(
+            input_dim=vocabulary_size,
+            output_dim=hidden_dim,
+            embeddings_initializer=esm2_kernel_initializer(),
+            dtype=dtype,
+            name="token_embedding",
+        )
+        if position_embedding_type == "absolute":
+            self.position_embedding = PositionEmbedding(
+                initializer=esm2_kernel_initializer(),
+                sequence_length=max_sequence_length,
+                dtype=dtype,
+                name="position_embedding",
+            )
+            self.embeddings_add = keras.layers.Add(
+                dtype=dtype,
+                name="embeddings_add",
+            )
+
+        self.output_layer_norm = keras.layers.LayerNormalization(
+            epsilon=layer_norm_eps,
+            dtype=dtype,
+            name="output_layer_norm",
+        )
+        if use_pre_layer_norm:
+            self.emb_layer_norm = keras.layers.LayerNormalization(
+                epsilon=layer_norm_eps,
+                dtype=dtype,
+                name="emb_layer_norm",
+            )
+        self.transformer_layers = []
+        for i in range(num_layers):
+            layer = ESMEncoder(
+                heads=num_heads,
+                head_size=head_size,
+                intermediate_size=intermediate_dim,
+                use_bias=use_bias,
+                max_wavelength=max_wavelength,
+                dropout=dropout,
+                activation=activation,
+                kernel_initializer=esm2_kernel_initializer(),
+                layer_norm_eps=layer_norm_eps,
+                dtype=dtype,
+                use_rotary=position_embedding_type == "rotary",
+                name=f"transformer_layer_{i}",
+            )
+            self.transformer_layers.append(layer)
+
+        # === Functional Model ===
+        token_id_input = keras.Input(
+            shape=(None,), dtype="int32", name="token_ids"
+        )
+
+        attention_mask = keras.ops.not_equal(token_id_input, pad_token_id)
+
+        token_vector = self.token_embedding(token_id_input)
+        if position_embedding_type == "absolute":
+            position_vector = self.position_embedding(
+                token_vector, start_index=pad_token_id
+            )
+            x = self.embeddings_add([token_vector, position_vector])
+        else:
+            x = token_vector
+        if use_pre_layer_norm:
+            x = self.emb_layer_norm(x)
+        for transformer_layer in self.transformer_layers:
+            x = transformer_layer(x, attention_mask=attention_mask)
+        output = self.output_layer_norm(x)
+        super().__init__(
+            inputs={
+                "token_ids": token_id_input,
+            },
+            outputs=output,
+            dtype=dtype,
+            **kwargs,
+        )
+
+        # === Config ===
+        self.vocabulary_size = vocabulary_size
+        self.num_layers = num_layers
+        self.num_heads = num_heads
+        self.hidden_dim = hidden_dim
+        self.intermediate_dim = intermediate_dim
+        self.dropout = dropout
+        self.max_wavelength = max_wavelength
+        self.head_size = head_size
+        self.activation = activations.get(activation)
+        self.use_bias = use_bias
+        self.start_token_index = 0
+        self.layer_norm_eps = layer_norm_eps
+        self.max_sequence_length = max_sequence_length
+        self.use_pre_layer_norm = use_pre_layer_norm
+        self.position_embedding_type = position_embedding_type
+        self.pad_token_id = pad_token_id
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "vocabulary_size": self.vocabulary_size,
+                "num_layers": self.num_layers,
+                "num_heads": self.num_heads,
+                "hidden_dim": self.hidden_dim,
+                "intermediate_dim": self.intermediate_dim,
+                "dropout": self.dropout,
+                "max_wavelength": self.max_wavelength,
+                "use_bias": self.use_bias,
+                "activation": activations.serialize(self.activation),
+                "layer_norm_eps": self.layer_norm_eps,
+                "use_pre_layer_norm": self.use_pre_layer_norm,
+                "position_embedding_type": self.position_embedding_type,
+                "max_sequence_length": self.max_sequence_length,
+                "pad_token_id": self.pad_token_id,
+            }
+        )
+        return config

keras_hub/src/models/esm/esm_classifier.py

@@ -0,0 +1,184 @@
+import keras
+
+from keras_hub.src.api_export import keras_hub_export
+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_classifier_preprocessor import (
+    ESMProteinClassifierPreprocessor,
+)
+from keras_hub.src.models.text_classifier import TextClassifier
+
+
+@keras_hub_export("keras_hub.models.ESMProteinClassifier")
+class ESMProteinClassifier(TextClassifier):
+    """An end-to-end ESM model for classification tasks.
+
+    This model attaches a classification head to
+    `keras_hub.models.ESMBackbone`, mapping from the backbone outputs
+    to logits suitable for a classification task. For usage of this model with
+    pre-trained weights, use the `from_preset()` constructor.
+
+    This model can optionally be configured with a `preprocessor` layer, in
+    which case it will automatically apply preprocessing to raw inputs during
+    `fit()`, `predict()`, and `evaluate()`. This is done by default when
+    creating the model with `from_preset()`.
+
+    Args:
+        backbone: A `keras_hub.models.ESMBackbone` instance.
+        num_classes: int. Number of classes to predict.
+        preprocessor: A `keras_hub.models.ESMProteinClassifierPreprocessor`
+            or `None`. If `None`, this model will not apply preprocessing, and
+            inputs should be preprocessed before calling the model.
+        activation: Optional `str` or callable. The
+            activation function to use on the model outputs. Set
+            `activation="softmax"` to return output probabilities.
+            Defaults to `None`.
+        dropout: float. The dropout probability value, applied after the dense
+            layer.
+
+    Examples:
+
+    Raw string data.
+    ```python
+    features = ["The quick brown fox jumped.", "I forgot my homework."]
+    labels = [0, 3]
+
+    # Pretrained classifier.
+    classifier = keras_hub.models.ESMProteinClassifier.from_preset(
+        hf://facebook/esm2_t6_8M_UR50D,
+        num_classes=4,
+    )
+    classifier.fit(x=features, y=labels, batch_size=2)
+    classifier.predict(x=features, batch_size=2)
+
+    # Re-compile (e.g., with a new learning rate).
+    classifier.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`).
+    classifier.backbone.trainable = False
+    # Fit again.
+    classifier.fit(x=features, y=labels, batch_size=2)
+    ```
+
+    Preprocessed integer data.
+    ```python
+    features = {
+        "token_ids": np.ones(shape=(2, 12), dtype="int32"),
+    }
+    labels = [0, 3]
+
+    # Pretrained classifier without preprocessing.
+    classifier = keras_hub.models.ESMProteinClassifier.from_preset(
+        hf://facebook/esm2_t6_8M_UR50D,
+        num_classes=4,
+        preprocessor=None,
+    )
+    classifier.fit(x=features, y=labels, batch_size=2)
+    ```
+
+    Custom backbone and vocabulary.
+    ```python
+    features = ["The quick brown fox jumped.", "I forgot my homework."]
+    labels = [0, 3]
+
+    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=tokenizer,
+        sequence_length=128,
+    )
+    backbone = keras_hub.models.ESMBackbone(
+        vocabulary_size=30552,
+        num_layers=4,
+        num_heads=4,
+        hidden_dim=256,
+        intermediate_dim=512,
+        max_wavelength=128,
+        num_head=4,
+    )
+    classifier = keras_hub.models.ESMProteinClassifier(
+        backbone=backbone,
+        preprocessor=preprocessor,
+        num_classes=4,
+    )
+    classifier.fit(x=features, y=labels, batch_size=2)
+    ```
+    """
+
+    backbone_cls = ESMBackbone
+    preprocessor_cls = ESMProteinClassifierPreprocessor
+
+    def __init__(
+        self,
+        backbone,
+        num_classes,
+        preprocessor=None,
+        activation=None,
+        hidden_dim=None,
+        dropout=0.0,
+        **kwargs,
+    ):
+        # === Layers ===
+        self.backbone = backbone
+        self.preprocessor = preprocessor
+        self.pooled_dropout = keras.layers.Dropout(
+            dropout,
+            dtype=backbone.dtype_policy,
+            name="pooled_dropout",
+        )
+        hidden_dim = hidden_dim or backbone.hidden_dim
+        self.pooled_dense = keras.layers.Dense(
+            hidden_dim,
+            activation="tanh",
+            dtype=backbone.dtype_policy,
+            name="pooled_dense",
+        )
+        self.output_dropout = keras.layers.Dropout(
+            dropout,
+            dtype=backbone.dtype_policy,
+            name="output_dropout",
+        )
+        self.output_dense = keras.layers.Dense(
+            num_classes,
+            kernel_initializer=esm2_kernel_initializer(),
+            activation=activation,
+            dtype=backbone.dtype_policy,
+            name="logits",
+        )
+
+        # === Functional Model ===
+        inputs = backbone.input
+        x = backbone(inputs)[:, backbone.start_token_index, :]
+        x = self.pooled_dropout(x)
+        x = self.pooled_dense(x)
+        x = self.output_dropout(x)
+        outputs = self.output_dense(x)
+        super().__init__(
+            inputs=inputs,
+            outputs=outputs,
+            **kwargs,
+        )
+
+        # === Config ===
+        self.num_classes = num_classes
+        self.activation = keras.activations.get(activation)
+        self.hidden_dim = hidden_dim
+        self.dropout = dropout
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "num_classes": self.num_classes,
+                "activation": keras.activations.serialize(self.activation),
+                "hidden_dim": self.hidden_dim,
+                "dropout": self.dropout,
+            }
+        )
+        return config

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)