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

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_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/roformer_v2/roformer_v2_attention.py

@@ -1,6 +1,7 @@
 import keras
 from keras import initializers
 from keras import ops
+from packaging import version
 
 
 class RoformerNorm(keras.layers.Layer):
@@ -179,7 +180,7 @@ class RoformerAttention(keras.layers.Layer):
         vw = ops.reshape(vw, (b, s, self.heads, self.head_size))
 
         qw, kw = self.rotary_embedding_layer([qw, kw])
-        if keras.__version__ < "3.6":
+        if version.parse(keras.__version__) < version.parse("3.6"):
             raise ("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])

keras_hub/src/utils/transformers/convert_esm.py

@@ -0,0 +1,159 @@
+import numpy as np
+
+from keras_hub.src.models.esm.esm_backbone import ESMBackbone
+from keras_hub.src.utils.preset_utils import get_file
+
+backbone_cls = ESMBackbone
+
+
+def convert_backbone_config(transformers_config):
+    return {
+        "vocabulary_size": transformers_config["vocab_size"],
+        "num_layers": transformers_config["num_hidden_layers"],
+        "num_heads": transformers_config["num_attention_heads"],
+        "hidden_dim": transformers_config["hidden_size"],
+        "intermediate_dim": transformers_config["intermediate_size"],
+        "dropout": transformers_config["hidden_dropout_prob"],
+        "position_embedding_type": transformers_config[
+            "position_embedding_type"
+        ],
+        "pad_token_id": transformers_config["pad_token_id"],
+        "max_sequence_length": transformers_config.get(
+            "max_position_embeddings", None
+        ),
+        "layer_norm_eps": transformers_config.get("layer_norm_eps", 1e-12),
+        "use_pre_layer_norm": transformers_config.get(
+            "emb_layer_norm_before", False
+        ),
+        "activation": transformers_config.get("activation", "gelu"),
+        "max_wavelength": transformers_config.get("max_wavelength", 10000),
+    }
+
+
+def transpose_and_reshape(x, shape):
+    return np.reshape(np.transpose(x), shape)
+
+
+def convert_weights(backbone, loader, transformers_config):
+    # Embedding layer
+    loader.port_weight(
+        keras_variable=backbone.get_layer("token_embedding").embeddings,
+        hf_weight_key="embeddings.word_embeddings.weight",
+    )
+    if transformers_config["position_embedding_type"] == "absolute":
+        loader.port_weight(
+            keras_variable=backbone.get_layer(
+                "position_embedding"
+            ).position_embeddings,
+            hf_weight_key="embeddings.position_embeddings.weight",
+        )
+    if transformers_config.get("emb_layer_norm_before", False):
+        loader.port_weight(
+            keras_variable=backbone.get_layer("emb_layer_norm").gamma,
+            hf_weight_key="embeddings.layer_norm.weight",
+        )
+        loader.port_weight(
+            keras_variable=backbone.get_layer("emb_layer_norm").beta,
+            hf_weight_key="embeddings.layer_norm.bias",
+        )
+
+    loader.port_weight(
+        keras_variable=backbone.output_layer_norm.gamma,
+        hf_weight_key="encoder.emb_layer_norm_after.weight",
+    )
+    loader.port_weight(
+        keras_variable=backbone.output_layer_norm.beta,
+        hf_weight_key="encoder.emb_layer_norm_after.bias",
+    )
+
+    # Attention blocks
+    for i in range(backbone.num_layers):
+        block = backbone.get_layer(f"transformer_layer_{i}")
+        attn = block.attention_layer
+        hf_prefix = "encoder.layer."
+        # Attention layers
+        loader.port_weight(
+            keras_variable=attn.q_dense.kernel,
+            hf_weight_key=f"{hf_prefix}{i}.attention.self.query.weight",
+            hook_fn=transpose_and_reshape,
+        )
+        loader.port_weight(
+            keras_variable=attn.q_dense.bias,
+            hf_weight_key=f"{hf_prefix}{i}.attention.self.query.bias",
+            hook_fn=lambda hf_tensor, shape: np.reshape(hf_tensor, shape),
+        )
+        loader.port_weight(
+            keras_variable=attn.k_dense.kernel,
+            hf_weight_key=f"{hf_prefix}{i}.attention.self.key.weight",
+            hook_fn=transpose_and_reshape,
+        )
+        loader.port_weight(
+            keras_variable=attn.k_dense.bias,
+            hf_weight_key=f"{hf_prefix}{i}.attention.self.key.bias",
+            hook_fn=lambda hf_tensor, shape: np.reshape(hf_tensor, shape),
+        )
+        loader.port_weight(
+            keras_variable=attn.v_dense.kernel,
+            hf_weight_key=f"{hf_prefix}{i}.attention.self.value.weight",
+            hook_fn=transpose_and_reshape,
+        )
+        loader.port_weight(
+            keras_variable=attn.v_dense.bias,
+            hf_weight_key=f"{hf_prefix}{i}.attention.self.value.bias",
+            hook_fn=lambda hf_tensor, shape: np.reshape(hf_tensor, shape),
+        )
+        loader.port_weight(
+            keras_variable=attn.o_dense.kernel,
+            hf_weight_key=f"{hf_prefix}{i}.attention.output.dense.weight",
+            hook_fn=transpose_and_reshape,
+        )
+        loader.port_weight(
+            keras_variable=attn.o_dense.bias,
+            hf_weight_key=f"{hf_prefix}{i}.attention.output.dense.bias",
+            hook_fn=lambda hf_tensor, shape: np.reshape(hf_tensor, shape),
+        )
+        # Attention layer norm.
+        loader.port_weight(
+            keras_variable=block.attention_norm.gamma,
+            hf_weight_key=f"{hf_prefix}{i}.attention.LayerNorm.weight",
+        )
+        loader.port_weight(
+            keras_variable=block.attention_norm.beta,
+            hf_weight_key=f"{hf_prefix}{i}.attention.LayerNorm.bias",
+        )
+        # MLP layers
+        loader.port_weight(
+            keras_variable=block.feedforward_intermediate_dense.kernel,
+            hf_weight_key=f"{hf_prefix}{i}.intermediate.dense.weight",
+            hook_fn=lambda hf_tensor, _: np.transpose(hf_tensor, axes=(1, 0)),
+        )
+        loader.port_weight(
+            keras_variable=block.feedforward_intermediate_dense.bias,
+            hf_weight_key=f"{hf_prefix}{i}.intermediate.dense.bias",
+        )
+        loader.port_weight(
+            keras_variable=block.feedforward_output_dense.kernel,
+            hf_weight_key=f"{hf_prefix}{i}.output.dense.weight",
+            hook_fn=lambda hf_tensor, _: np.transpose(hf_tensor, axes=(1, 0)),
+        )
+        loader.port_weight(
+            keras_variable=block.feedforward_output_dense.bias,
+            hf_weight_key=f"{hf_prefix}{i}.output.dense.bias",
+        )
+        # Output layer norm.
+        loader.port_weight(
+            keras_variable=block.feedforward_norm.gamma,
+            hf_weight_key=f"{hf_prefix}{i}.LayerNorm.weight",
+        )
+        loader.port_weight(
+            keras_variable=block.feedforward_norm.beta,
+            hf_weight_key=f"{hf_prefix}{i}.LayerNorm.bias",
+        )
+
+
+def convert_tokenizer(cls, preset, **kwargs):
+    return cls(
+        get_file(preset, "vocab.txt"),
+        lowercase=False,
+        **kwargs,
+    )

keras_hub/src/utils/transformers/preset_loader.py

@@ -9,6 +9,7 @@ from keras_hub.src.utils.transformers import convert_bert
 from keras_hub.src.utils.transformers import convert_deit
 from keras_hub.src.utils.transformers import convert_dinov2
 from keras_hub.src.utils.transformers import convert_distilbert
+from keras_hub.src.utils.transformers import convert_esm
 from keras_hub.src.utils.transformers import convert_gemma
 from keras_hub.src.utils.transformers import convert_gpt2
 from keras_hub.src.utils.transformers import convert_llama3
@@ -38,6 +39,8 @@ class TransformersPresetLoader(PresetLoader):
             self.converter = convert_distilbert
         elif model_type in ("dinov2", "dinov2_with_registers"):
             self.converter = convert_dinov2
+        elif model_type == "esm":
+            self.converter = convert_esm
         elif model_type in ("gemma", "gemma2"):
             self.converter = convert_gemma
         elif model_type == "gpt2":

keras_hub/src/version.py

@@ -1,7 +1,7 @@
 from keras_hub.src.api_export import keras_hub_export
 
 # Unique source of truth for the version number.
-__version__ = "0.22.0.dev202508100425"
+__version__ = "0.22.0.dev202508120417"
 
 
 @keras_hub_export("keras_hub.version")