feat: update backbone + add tokenizer

SauravMaheshkar · SauravMaheshkar · commit 5b6e62aa0644 · 2025-05-17T04:01:18.000+01:00
diff --git a/keras_hub/src/models/modernbert/modernbert_backbone.py b/keras_hub/src/models/modernbert/modernbert_backbone.py
@@ -5,20 +5,15 @@
     ReversibleEmbedding,
 )
 from keras_hub.src.layers.modeling.rotary_embedding import RotaryEmbedding
-from keras_hub.src.layers.modeling.transformer_encoder import TransformerEncoder
 from keras_hub.src.models.backbone import Backbone
-from keras_hub.src.models.gemma.rms_normalization import RMSNormalization
+from keras_hub.src.models.modernbert.modernbert_layers import (
+    ModernBERTEncoderLayer,
+)
 from keras_hub.src.utils.keras_utils import gelu_approximate
 
 
 @keras_hub_export("keras_hub.models.ModernBertBackbone")
 class ModernBertBackbone(Backbone):
-    """A ModernBERT encoder network.
-
-    This class implements the ModernBERT backbone, using rotary embeddings,
-    RMS normalization, and a stack of TransformerEncoder layers.
-    """
-
     def __init__(
         self,
         vocabulary_size,
@@ -45,37 +40,33 @@ def __init__(
         )
         self.position_embedding = RotaryEmbedding(
             max_wavelength=rotary_max_wavelength,
-            sequence_axis=1,
-            feature_axis=-1,
             dtype=dtype,
             name="rotary_embedding",
         )
-        self.embeddings_layer_norm = RMSNormalization(
-            dtype=dtype,
+        self.embeddings_layer_norm = keras.layers.LayerNormalization(
             epsilon=layer_norm_epsilon,
-        )
-        self.embeddings_dropout = keras.layers.Dropout(
-            dropout, dtype=dtype, name="embeddings_dropout"
+            dtype=dtype,
+            rms_scaling=True,
+            name="embeddings_layer_norm",
         )
         self.transformer_layers = []
         for i in range(num_layers):
-            layer = TransformerEncoder(
+            layer = ModernBERTEncoderLayer(
+                hidden_size=hidden_dim,
+                intermediate_size=intermediate_dim,
                 num_heads=num_heads,
-                intermediate_dim=intermediate_dim,
                 activation=gelu_approximate,
-                dropout=dropout,
                 layer_norm_epsilon=layer_norm_epsilon,
-                kernel_initializer=keras.initializers.TruncatedNormal(
-                    stddev=0.02
-                ),
+                rotary_embedding=self.position_embedding,
                 dtype=dtype,
                 name=f"transformer_layer_{i}",
             )
             self.transformer_layers.append(layer)
-        self.final_norm = RMSNormalization(
-            dtype=dtype,
+        self.final_norm = keras.layers.LayerNormalization(
             epsilon=layer_norm_epsilon,
-            name="final_normalization",
+            rms_scaling=True,
+            dtype=dtype,
+            name="final_layernorm",
         )
 
         # === Functional Model ===
@@ -85,20 +76,13 @@ def __init__(
         padding_mask_input = keras.Input(
             shape=(None,), dtype="int32", name="padding_mask"
         )
-
-        # Embed tokens and apply rotary position embedding
         x = self.token_embedding(token_id_input)
-        x = self.position_embedding(x)
         x = self.embeddings_layer_norm(x)
-        x = self.embeddings_dropout(x)
-
-        # Transformer layers
         for transformer_layer in self.transformer_layers:
-            x = transformer_layer(x, padding_mask=padding_mask_input)
-
-        # Final normalization
+            x = transformer_layer(x)
         sequence_output = self.final_norm(x)
 
+        # Instantiate using Functional API Model constructor
         super().__init__(
             inputs={
                 "token_ids": token_id_input,
diff --git a/keras_hub/src/models/modernbert/modernbert_layers.py b/keras_hub/src/models/modernbert/modernbert_layers.py
@@ -0,0 +1,95 @@
+import keras
+from keras import layers
+from keras import ops
+
+from keras_hub.src.models.flux.flux_maths import rearrange_symbolic_tensors
+from keras_hub.src.models.flux.flux_maths import scaled_dot_product_attention
+
+
+class MLP(keras.layers.Layer):
+    def __init__(
+        self,
+        hidden_size,
+        intermediate_size,
+        activation="gelu",
+        dtype=None,
+        **kwargs,
+    ):
+        super(MLP, self).__init__(**kwargs)
+        self.Wi = layers.Dense(
+            intermediate_size * 2,
+            use_bias=False,
+            dtype=dtype,
+        )
+        self.act = keras.activations.get(activation)
+        self.Wo = layers.Dense(
+            hidden_size,
+            use_bias=False,
+            dtype=dtype,
+        )
+
+    def call(self, x):
+        input, gate = ops.split(self.Wi(x), 2, axis=-1)
+        return self.Wo(self.act(input) * gate)
+
+
+class ModernBERTAttention(keras.Model):
+    def __init__(
+        self, hidden_size, num_heads, rotary_embedding, dtype=None, **kwargs
+    ):
+        super(ModernBERTAttention, self).__init__(**kwargs)
+        self.num_heads = num_heads
+        self.hidden_size = hidden_size
+        self.rotary_embedding = rotary_embedding
+        self.Wqkv = layers.Dense(hidden_size * 3, use_bias=False, dtype=dtype)
+        self.Wo = layers.Dense(hidden_size, use_bias=False, dtype=dtype)
+
+    def build(self, input_shape):
+        self.Wqkv.build(input_shape)
+        self.Wo.build((None, input_shape[1], input_shape[-1]))
+
+    def call(self, x):
+        qkv = self.Wqkv(x)
+        q, k, v = rearrange_symbolic_tensors(qkv, K=3, H=self.num_heads)
+
+        # Apply rotary embeddings
+        q = self.rotary_embedding(q)
+        k = self.rotary_embedding(k)
+
+        # Apply scaled dot product attention
+        x = scaled_dot_product_attention(q, k, v)
+
+        # Reshape and apply final dense layer
+        x = ops.transpose(x, (0, 2, 1, 3))
+        b, s, h, d = ops.shape(x)
+        x = ops.reshape(x, (b, s, h * d))
+        x = self.Wo(x)
+        return x
+
+
+class ModernBERTEncoderLayer(keras.Model):
+    def __init__(
+        self,
+        hidden_size,
+        intermediate_size,
+        num_heads,
+        activation="gelu",
+        layer_norm_epsilon=1e-05,
+        rotary_embedding=None,
+        dtype=None,
+        **kwargs,
+    ):
+        super(ModernBERTEncoderLayer, self).__init__(**kwargs)
+        self.attn = ModernBERTAttention(
+            hidden_size, num_heads, rotary_embedding, dtype=dtype
+        )
+        self.mlp_norm = layers.LayerNormalization(
+            epsilon=layer_norm_epsilon, dtype=dtype
+        )
+        self.mlp = MLP(hidden_size, intermediate_size, activation, dtype=dtype)
+
+    def call(self, x):
+        x = self.attn(x)
+        x = self.mlp_norm(x)
+        x = self.mlp(x)
+        return x
diff --git a/keras_hub/src/models/modernbert/modernbert_tokenizer.py b/keras_hub/src/models/modernbert/modernbert_tokenizer.py
@@ -0,0 +1,36 @@
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.models.modernbert.modernbert_backbone import (
+    ModernBertBackbone,
+)
+from keras_hub.src.tokenizers.byte_pair_tokenizer import BytePairTokenizer
+
+
+@keras_hub_export(
+    [
+        "keras_hub.tokenizers.ModernBertTokenizer",
+        "keras_hub.models.ModernBertTokenizer",
+    ]
+)
+class ModernBertTokenizer(BytePairTokenizer):
+    backbone_cls = ModernBertBackbone
+
+    def __init__(
+        self,
+        vocabulary=None,
+        merges=None,
+        **kwargs,
+    ):
+        self._add_special_token("[CLS]", "cls_token")
+        self._add_special_token("[SEP]", "sep_token")
+        self._add_special_token("[PAD]", "pad_token")
+        self._add_special_token("[UNK]", "unk_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("[SEP]", "end_token")
+        super().__init__(
+            vocabulary=vocabulary,
+            merges=merges,
+            **kwargs,
+        )
diff --git a/keras_hub/src/models/modernbert/modernbert_tokenizer_test.py b/keras_hub/src/models/modernbert/modernbert_tokenizer_test.py
@@ -0,0 +1,35 @@
+from keras_hub.src.models.modernbert.modernbert_tokenizer import (
+    ModernBertTokenizer,
+)
+from keras_hub.src.tests.test_case import TestCase
+
+
+class ModernBertTokenizerTest(TestCase):
+    def setUp(self):
+        self.vocab = ["[CLS]", "[PAD]", "[SEP]", "air", "Ġair", "plane", "Ġat"]
+        self.vocab += ["port", "[MASK]", "[UNK]"]
+        self.vocab = dict([(token, i) for i, token in enumerate(self.vocab)])
+        self.merges = ["Ġ a", "Ġ t", "Ġ i", "Ġ b", "a i", "p l", "n e"]
+        self.merges += ["Ġa t", "p o", "r t", "Ġt h", "ai r", "pl a", "po rt"]
+        self.merges += ["Ġai r", "Ġa i", "pla ne"]
+        self.init_kwargs = {"vocabulary": self.vocab, "merges": self.merges}
+        self.input_data = [
+            "[CLS] airplane at airport[SEP][PAD]",
+            " airplane airport",
+        ]
+
+    def test_tokenizer_basics(self):
+        self.run_preprocessing_layer_test(
+            cls=ModernBertTokenizer,
+            init_kwargs=self.init_kwargs,
+            input_data=self.input_data,
+            expected_output=[[0, 4, 5, 6, 4, 7, 2, 1], [4, 5, 4, 7]],
+            expected_detokenize_output=[
+                "[CLS] airplane at airport[SEP][PAD]",
+                " airplane airport",
+            ],
+        )
+
+    def test_errors_missing_special_tokens(self):
+        with self.assertRaises(ValueError):
+            ModernBertTokenizer(vocabulary=["a", "b", "c"], merges=[])
diff --git a/tools/checkpoint_conversion/convert_modernbert_checkpoints.py b/tools/checkpoint_conversion/convert_modernbert_checkpoints.py
@@ -0,0 +1,120 @@
+"""Convert ModernBERT checkpoints.
+
+python tools/checkpoint_conversion/convert_modernbert_checkpoints.py \
+    --preset modernbert_base
+python tools/checkpoint_conversion/convert_modernbert_checkpoints.py \
+    --preset modernbert_large
+"""
+
+import json
+import os
+
+import numpy as np
+import requests
+import transformers
+from absl import app
+from absl import flags
+
+from keras_hub.src.models.modernbert.modernbert_backbone import (
+    ModernBertBackbone,
+)
+
+PRESET_MAP = {
+    "modernbert_base": "answerdotai/ModernBERT-base",
+    "modernbert_large": "answerdotai/ModernBERT-large",
+}
+
+EXTRACT_DIR = "./{}"
+
+FLAGS = flags.FLAGS
+flags.DEFINE_string(
+    "preset",
+    None,
+    f"Must be one of {','.join(PRESET_MAP.keys())}",
+)
+
+
+def download_files(hf_model_name):
+    extract_dir = EXTRACT_DIR.format(FLAGS.preset)
+    if not os.path.exists(extract_dir):
+        os.makedirs(extract_dir)
+
+    # Config.
+    config_path = os.path.join(extract_dir, "config.json")
+    response = requests.get(
+        f"https://huggingface.co/{hf_model_name}/raw/main/config.json"
+    )
+    open(config_path, "wb").write(response.content)
+
+
+def convert_model(hf_model):
+    extract_dir = EXTRACT_DIR.format(FLAGS.preset)
+    config_path = os.path.join(extract_dir, "config.json")
+
+    # Build config.
+    cfg = {}
+    with open(config_path, "r") as pt_cfg_handler:
+        pt_cfg = json.load(pt_cfg_handler)
+    cfg["vocabulary_size"] = pt_cfg["vocab_size"]
+    cfg["num_layers"] = pt_cfg["num_hidden_layers"]
+    cfg["num_heads"] = pt_cfg["num_attention_heads"]
+    cfg["hidden_dim"] = pt_cfg["hidden_size"]
+    cfg["intermediate_dim"] = pt_cfg["intermediate_size"]
+    cfg["dropout"] = pt_cfg["embedding_dropout"]
+    cfg["max_sequence_length"] = pt_cfg["max_position_embeddings"]
+
+    return ModernBertBackbone(**cfg)
+
+
+def convert_weights(keras_model, hf_model):
+    # Get `state_dict` from `hf_model`.
+    state_dict = hf_model.state_dict()
+
+    keras_model.get_layer("token_embedding").set_weights(
+        [np.asarray(state_dict["embeddings.tok_embeddings.weight"])]
+    )
+
+    keras_model.get_layer("embeddings_layer_norm").set_weights(
+        [np.asarray(state_dict["embeddings.norm.weight"])]
+    )
+
+    for i in range(keras_model.num_layers):
+        keras_model.transformer_layers[i].attn.Wqkv.kernel.assign(
+            state_dict[f"layers.{i}.attn.Wqkv.weight"].T
+        )
+        keras_model.transformer_layers[i].attn.Wo.kernel.assign(
+            state_dict[f"layers.{i}.attn.Wo.weight"]
+        )
+        keras_model.transformer_layers[i].mlp_norm.gamma.assign(
+            state_dict[f"layers.{i}.mlp_norm.weight"]
+        )
+        keras_model.transformer_layers[i].mlp.Wi.kernel.assign(
+            state_dict[f"layers.{i}.mlp.Wi.weight"].T
+        )
+        keras_model.transformer_layers[i].mlp.Wo.kernel.assign(
+            state_dict[f"layers.{i}.mlp.Wo.weight"].T
+        )
+
+    keras_model.get_layer("final_layernorm").set_weights(
+        [np.asarray(state_dict["final_norm.weight"])]
+    )
+
+
+def main(_):
+    hf_model_name = PRESET_MAP[FLAGS.preset]
+    download_files(hf_model_name)
+
+    hf_model = transformers.AutoModel.from_pretrained(hf_model_name)
+    hf_model.eval()
+
+    print(f"🏃 Coverting {FLAGS.preset}")
+    keras_model = convert_model(hf_model)
+    print("✅ KerasHub model loaded.")
+
+    convert_weights(keras_model, hf_model)
+    print("✅ Weights converted.")
+
+
+if __name__ == "__main__":
+    flags.mark_flag_as_required("preset")
+    app.run(main)
