keras-hub-nightly 0.20.0.dev202504010407__tar.gz → 0.20.0.dev202504030357__tar.gz

{keras_hub_nightly-0.20.0.dev202504010407 → keras_hub_nightly-0.20.0.dev202504030357}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: keras-hub-nightly
-Version: 0.20.0.dev202504010407
+Version: 0.20.0.dev202504030357
 Summary: Industry-strength Natural Language Processing extensions for Keras.
 Home-page: https://github.com/keras-team/keras-hub
 Author: Keras team

{keras_hub_nightly-0.20.0.dev202504010407 → keras_hub_nightly-0.20.0.dev202504030357}/keras_hub/api/models/__init__.py

@@ -323,6 +323,24 @@ from keras_hub.src.models.roberta.roberta_text_classifier_preprocessor import (
     RobertaTextClassifierPreprocessor as RobertaPreprocessor,
 )
 from keras_hub.src.models.roberta.roberta_tokenizer import RobertaTokenizer
+from keras_hub.src.models.roformer_v2.roformer_v2_backbone import (
+    RoformerV2Backbone as RorformerV2Backbone,
+)
+from keras_hub.src.models.roformer_v2.roformer_v2_masked_lm import (
+    RoformerV2MaskedLM,
+)
+from keras_hub.src.models.roformer_v2.roformer_v2_masked_lm_preprocessor import (
+    RoformerV2MaskedLMPreprocessor,
+)
+from keras_hub.src.models.roformer_v2.roformer_v2_text_classifier import (
+    RorformerV2TextClassifier,
+)
+from keras_hub.src.models.roformer_v2.roformer_v2_text_classifier_preprocessor import (
+    RoformerV2TextClassifierPreprocessor,
+)
+from keras_hub.src.models.roformer_v2.roformer_v2_tokenizer import (
+    RoformerV2Tokenizer,
+)
 from keras_hub.src.models.sam.sam_backbone import SAMBackbone
 from keras_hub.src.models.sam.sam_image_segmenter import SAMImageSegmenter
 from keras_hub.src.models.sam.sam_image_segmenter_preprocessor import (

{keras_hub_nightly-0.20.0.dev202504010407 → keras_hub_nightly-0.20.0.dev202504030357}/keras_hub/api/tokenizers/__init__.py

@@ -35,6 +35,9 @@ from keras_hub.src.models.qwen.qwen_tokenizer import (
     QwenTokenizer as Qwen2Tokenizer,
 )
 from keras_hub.src.models.roberta.roberta_tokenizer import RobertaTokenizer
+from keras_hub.src.models.roformer_v2.roformer_v2_tokenizer import (
+    RoformerV2Tokenizer,
+)
 from keras_hub.src.models.siglip.siglip_tokenizer import SigLIPTokenizer
 from keras_hub.src.models.t5.t5_tokenizer import T5Tokenizer
 from keras_hub.src.models.whisper.whisper_tokenizer import WhisperTokenizer

{keras_hub_nightly-0.20.0.dev202504010407 → keras_hub_nightly-0.20.0.dev202504030357}/keras_hub/src/models/gemma/gemma_attention.py

@@ -6,7 +6,9 @@ from keras import ops
 
 from keras_hub.src.layers.modeling.rotary_embedding import RotaryEmbedding
 from keras_hub.src.utils.keras_utils import clone_initializer
-from keras_hub.src.utils.keras_utils import has_flash_attention_support
+from keras_hub.src.utils.keras_utils import fused_attention_op_available
+from keras_hub.src.utils.keras_utils import gpu_supports_fused_attention_op
+from keras_hub.src.utils.keras_utils import running_on_gpu
 from keras_hub.src.utils.keras_utils import running_on_tpu
 
 
@@ -106,17 +108,22 @@ class CachedGemmaAttention(keras.layers.Layer):
         )
         return x
 
-    def _can_use_flash_attention(self):
-        if not has_flash_attention_support():
+    def _use_fused_attention_op(self):
+        if not fused_attention_op_available():
             return False
         if self.dropout > 0.0:
             return False
-        if self.logit_soft_cap is None:
-            return True
-        sig = inspect.signature(ops.dot_product_attention)
-        # We can currently only run soft capped attention for keras >= 3.10
-        # and only on TPU.
-        return running_on_tpu() and "attn_logits_soft_cap" in sig.parameters
+        if running_on_gpu():
+            # GPU never supports softcap in the fused op.
+            if self.logit_soft_cap is not None:
+                return False
+            return gpu_supports_fused_attention_op()
+        elif running_on_tpu():
+            # TPU supports softcap with on keras >= 3.10.
+            sig = inspect.signature(ops.dot_product_attention)
+            return "attn_logits_soft_cap" in sig.parameters
+        else:
+            return False
 
     def _compute_attention(
         self,
@@ -133,7 +140,14 @@ class CachedGemmaAttention(keras.layers.Layer):
             query_normalization = 1 / np.sqrt(
                 self.hidden_dim // self.num_query_heads
             )
-        if self._can_use_flash_attention():
+
+        if self.use_sliding_window_attention and attention_mask is not None:
+            attention_mask = self._mask_sliding_window(
+                attention_mask,
+                cache_update_index=cache_update_index,
+            )
+
+        if self._use_fused_attention_op():
             if attention_mask is not None:
                 attention_mask = ops.expand_dims(attention_mask, axis=1)
                 attention_mask = ops.cast(attention_mask, dtype="bool")
@@ -172,13 +186,8 @@ class CachedGemmaAttention(keras.layers.Layer):
                 ops.tanh(attention_logits), self.logit_soft_cap
             )
 
-        if self.use_sliding_window_attention:
-            attention_mask = self._mask_sliding_window(
-                attention_mask,
-                cache_update_index=cache_update_index,
-            )
-
-        attention_mask = attention_mask[:, None, None, :, :]
+        if attention_mask is not None:
+            attention_mask = attention_mask[:, None, None, :, :]
         orig_dtype = attention_logits.dtype
         attention_softmax = self.softmax(attention_logits, mask=attention_mask)
         attention_softmax = ops.cast(attention_softmax, orig_dtype)

{keras_hub_nightly-0.20.0.dev202504010407 → keras_hub_nightly-0.20.0.dev202504030357}/keras_hub/src/models/gemma3/gemma3_attention.py

@@ -7,7 +7,7 @@ from keras import ops
 from keras_hub.src.layers.modeling.rotary_embedding import RotaryEmbedding
 from keras_hub.src.models.gemma.rms_normalization import RMSNormalization
 from keras_hub.src.utils.keras_utils import clone_initializer
-from keras_hub.src.utils.keras_utils import has_flash_attention_support
+from keras_hub.src.utils.keras_utils import fused_attention_op_available
 from keras_hub.src.utils.keras_utils import running_on_tpu
 
 
@@ -140,7 +140,7 @@ class CachedGemma3Attention(keras.layers.Layer):
         return x
 
     def _can_use_flash_attention(self):
-        if not has_flash_attention_support():
+        if not fused_attention_op_available():
             return False
         if self.dropout > 0.0:
             return False

{keras_hub_nightly-0.20.0.dev202504010407 → keras_hub_nightly-0.20.0.dev202504030357}/keras_hub/src/models/gpt_neo_x/gpt_neo_x_attention.py

@@ -5,7 +5,7 @@ from keras import ops
 
 from keras_hub.src.layers.modeling.rotary_embedding import RotaryEmbedding
 from keras_hub.src.utils.keras_utils import clone_initializer
-from keras_hub.src.utils.keras_utils import has_flash_attention_support
+from keras_hub.src.utils.keras_utils import fused_attention_op_available
 
 
 class GPTNeoXAttention(keras.layers.Layer):
@@ -125,7 +125,7 @@ class GPTNeoXAttention(keras.layers.Layer):
     def _compute_attention(
         self, query, key, value, attention_mask=None, training=None
     ):
-        if has_flash_attention_support() and self.dropout == 0:
+        if fused_attention_op_available() and self.dropout == 0:
             # Use `dot_product_attention` with Flash Attention support if
             # available.
             if attention_mask is not None:

{keras_hub_nightly-0.20.0.dev202504010407 → keras_hub_nightly-0.20.0.dev202504030357}/keras_hub/src/models/llama/llama_attention.py

@@ -5,7 +5,7 @@ from keras import ops
 
 from keras_hub.src.layers.modeling.rotary_embedding import RotaryEmbedding
 from keras_hub.src.utils.keras_utils import clone_initializer
-from keras_hub.src.utils.keras_utils import has_flash_attention_support
+from keras_hub.src.utils.keras_utils import fused_attention_op_available
 
 
 class LlamaAttention(keras.layers.Layer):
@@ -185,7 +185,7 @@ class LlamaAttention(keras.layers.Layer):
         return self._softmax(attention_scores)
 
     def _compute_attention(self, query, key, value, attention_mask=None):
-        if has_flash_attention_support():
+        if fused_attention_op_available():
             # Use `dot_product_attention` with Flash Attention support if
             # available.
             if attention_mask is not None:

{keras_hub_nightly-0.20.0.dev202504010407 → keras_hub_nightly-0.20.0.dev202504030357}/keras_hub/src/models/mistral/mistral_attention.py

@@ -5,7 +5,7 @@ from keras import ops
 
 from keras_hub.src.layers.modeling.rotary_embedding import RotaryEmbedding
 from keras_hub.src.utils.keras_utils import clone_initializer
-from keras_hub.src.utils.keras_utils import has_flash_attention_support
+from keras_hub.src.utils.keras_utils import fused_attention_op_available
 
 
 # This is just a self-attention layer in Mistral. But it can be generalized
@@ -196,7 +196,7 @@ class CachedMistralAttention(keras.layers.Layer):
         return self._softmax(attention_scores)
 
     def _compute_attention(self, query, key, value, attention_mask=None):
-        if has_flash_attention_support():
+        if fused_attention_op_available():
             # Use `dot_product_attention` with Flash Attention support if
             # available.
             if attention_mask is not None:

{keras_hub_nightly-0.20.0.dev202504010407 → keras_hub_nightly-0.20.0.dev202504030357}/keras_hub/src/models/phi3/phi3_attention.py

@@ -8,7 +8,7 @@ from keras_hub.src.models.phi3.phi3_rotary_embedding import (
     Phi3SuScaledRotaryEmbedding,
 )
 from keras_hub.src.utils.keras_utils import clone_initializer
-from keras_hub.src.utils.keras_utils import has_flash_attention_support
+from keras_hub.src.utils.keras_utils import fused_attention_op_available
 
 
 class Phi3Attention(keras.layers.Layer):
@@ -217,7 +217,7 @@ class Phi3Attention(keras.layers.Layer):
         return self.softmax(attention_scores)
 
     def _compute_attention(self, query, key, value, attention_mask=None):
-        if has_flash_attention_support():
+        if fused_attention_op_available():
             # Use `dot_product_attention` with Flash Attention support if
             # available.
             if attention_mask is not None:

{keras_hub_nightly-0.20.0.dev202504010407 → keras_hub_nightly-0.20.0.dev202504030357}/keras_hub/src/models/qwen/qwen_attention.py

@@ -5,7 +5,7 @@ from keras import ops
 
 from keras_hub.src.layers.modeling.rotary_embedding import RotaryEmbedding
 from keras_hub.src.utils.keras_utils import clone_initializer
-from keras_hub.src.utils.keras_utils import has_flash_attention_support
+from keras_hub.src.utils.keras_utils import fused_attention_op_available
 
 
 class QwenAttention(keras.layers.Layer):
@@ -263,7 +263,7 @@ class QwenAttention(keras.layers.Layer):
         Returns:
             attention_output: Output tensor after applying attention.
         """
-        if has_flash_attention_support():
+        if fused_attention_op_available():
             # Use `dot_product_attention` with Flash Attention support if
             # available.
             if attention_mask is not None:

keras_hub_nightly-0.20.0.dev202504030357/keras_hub/src/models/roformer_v2/roformer_v2_attention.py

@@ -0,0 +1,212 @@
+import keras
+from keras import initializers
+from keras import ops
+
+
+class RoformerNorm(keras.layers.Layer):
+    """A normalization layer for Roformer that implements RMS normalization."""
+
+    def __init__(self, epsilon=1e-6, **kwargs):
+        super().__init__(**kwargs)
+        self.epsilon = epsilon
+
+    def build(self, input_shape):
+        dim = input_shape[-1]
+        self.scale = self.add_weight(
+            name="scale",
+            trainable=True,
+            shape=(dim,),
+            initializer="ones",
+            dtype=self.variable_dtype,
+        )
+        self.built = True
+
+    def call(self, x):
+        x = ops.cast(x, "float32")
+        var = ops.mean(ops.power(x, 2), axis=-1, keepdims=True)
+        x = x * ops.rsqrt(var + self.epsilon)
+        return ops.cast(x * self.scale, self.compute_dtype)
+
+    def get_config(self):
+        config = super().get_config()
+        config.update({"epsilon": self.epsilon})
+        return config
+
+
+class RoformrPositionalEmbedding(keras.layers.Layer):
+    """Native rotary implement by jianlin su
+    from native implement
+    https://github.com/bojone/bert4keras
+
+    """
+
+    def __init__(self, output_dim, max_wavelength=10000, **kwargs):
+        super().__init__(**kwargs)
+        self.max_wavelength = max_wavelength
+        self.output_dim = output_dim
+
+    def call(self, tensors):
+        input_shape = ops.shape(tensors[0])
+        seq_len = input_shape[1]
+        position_ids = ops.arange(0, seq_len, dtype=tensors[0].dtype)[None]
+        embeddings = self.sinusoidal_embeddings(
+            position_ids, self.output_dim, self.max_wavelength
+        )
+        embeddings = ops.cast(embeddings, self.compute_dtype)
+
+        ndim = ops.ndim(tensors[0])
+        sinusoidal = self.align(embeddings, [0, 1, -1], ndim)
+        cos_pos = ops.repeat(sinusoidal[..., 1::2], 2, -1)
+        sin_pos = ops.repeat(sinusoidal[..., ::2], 2, -1)
+        outputs = []
+        for tensor in tensors:
+            tensor2 = ops.stack([-tensor[..., 1::2], tensor[..., ::2]], ndim)
+            tensor2 = ops.reshape(tensor2, ops.shape(tensor))
+            outputs.append(tensor * cos_pos + tensor2 * sin_pos)
+        return outputs[0] if len(outputs) == 1 else outputs
+
+    def align(self, tensor, axes, ndim=None):
+        ndim = ndim or max(axes) + 1
+        indices = [None] * ndim
+        for i in axes:
+            indices[i] = slice(None)
+        if keras.config.backend() == "jax":
+            return tensor[tuple(indices)]
+        return tensor[indices]
+
+    def sinusoidal_embeddings(self, pos, dim, base=10000):
+        if dim % 2 != 0:
+            raise ("Dimension must be even")
+
+        indices = ops.arange(0, dim // 2, dtype="float32")
+        indices = ops.power(ops.cast(base, dtype="float32"), -2 * indices / dim)
+        embeddings = ops.einsum("...,d->...d", pos, indices)
+        embeddings = ops.stack(
+            [ops.sin(embeddings), ops.cos(embeddings)], axis=-1
+        )
+        shape = list(ops.shape(embeddings))
+        embeddings = ops.reshape(embeddings, shape[:-2] + [-1])
+        return embeddings
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "out_dim": self.out_dim,
+                "max_wavelength": self.max_wavelength,
+            }
+        )
+        return config
+
+
+@keras.saving.register_keras_serializable(package="keras_hub")
+class RoformerAttention(keras.layers.Layer):
+    """MultiHeadAttention by roformerV2
+
+    modifity from native implement
+    https://github.com/bojone/bert4keras
+    """
+
+    def __init__(
+        self,
+        heads,
+        head_size,
+        out_dim=None,
+        use_bias=False,
+        max_wavelength=10000,
+        kernel_initializer="glorot_uniform",
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.heads = heads
+        self.head_size = head_size
+        self.out_dim = out_dim or heads * head_size
+        self.use_bias = use_bias
+        self.kernel_initializer = initializers.get(kernel_initializer)
+        self.max_wavelength = max_wavelength
+
+    def build(self, input_shape):
+        super().build(input_shape)
+        self.q_dense = keras.layers.Dense(
+            units=self.head_size * self.heads,
+            use_bias=self.use_bias,
+            kernel_initializer=self.kernel_initializer,
+            name="q_dense_layer",
+            dtype=self.dtype_policy,
+        )
+        self.q_dense.build(input_shape)
+
+        self.k_dense = keras.layers.Dense(
+            units=self.head_size * self.heads,
+            use_bias=self.use_bias,
+            kernel_initializer=self.kernel_initializer,
+            name="k_dense_layer",
+            dtype=self.dtype_policy,
+        )
+        self.k_dense.build(input_shape)
+
+        self.v_dense = keras.layers.Dense(
+            units=self.head_size * self.heads,
+            use_bias=self.use_bias,
+            kernel_initializer=self.kernel_initializer,
+            name="v_dense_layer",
+            dtype=self.dtype_policy,
+        )
+        self.v_dense.build(input_shape)
+
+        self.o_dense = keras.layers.Dense(
+            units=self.out_dim,
+            use_bias=self.use_bias,
+            kernel_initializer=self.kernel_initializer,
+            name="o_dense_layer",
+            dtype=self.dtype_policy,
+        )
+        self.o_dense.build([None, None, self.head_size * self.heads])
+
+        self.rotary_embedding_layer = RoformrPositionalEmbedding(
+            self.head_size, 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))
+
+        qw, kw = self.rotary_embedding_layer([qw, kw])
+        if keras.__version__ < "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])
+        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 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,
+                "out_dim": self.out_dim,
+                "use_bias": self.use_bias,
+                "max_wavelength": self.max_wavelength,
+                "kernel_initializer": initializers.serialize(
+                    self.kernel_initializer
+                ),
+            }
+        )
+        return config

keras_hub_nightly-0.20.0.dev202504030357/keras_hub/src/models/roformer_v2/roformer_v2_backbone.py

@@ -0,0 +1,198 @@
+import keras
+from keras import activations
+
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.layers.modeling.reversible_embedding import (
+    ReversibleEmbedding,
+)
+from keras_hub.src.models.backbone import Backbone
+from keras_hub.src.models.roformer_v2.roformer_v2_attention import RoformerNorm
+from keras_hub.src.models.roformer_v2.roformer_v2_encoder import (
+    RoformerV2Encoder,
+)
+
+
+def roformer_kernel_initializer(stddev=0.02):
+    return keras.initializers.TruncatedNormal(stddev=stddev)
+
+
+@keras_hub_export("keras_hub.models.RorformerV2Backbone")
+class RoformerV2Backbone(Backbone):
+    """A RoformerV2 encoder network.
+
+    This class implements a bi-directional Transformer-based encoder as
+    described in ["Roformer"](https://github.com/ZhuiyiTechnology/roformer).
+    It includes the
+    embedding lookups and transformer layers, but not the masked language model
+    or next sentence prediction heads.
+
+    The default constructor gives a fully customizable, randomly initialized
+    RoformerV2 encoder with any number of layers, heads, and embed dim.To
+    load preset architectures and weights, use the `from_preset()` constructor.
+
+    Disclaimer: Pre-trained models are provided on an "as is" basis, without
+    warranties or conditions of any kind.
+
+    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.
+        num_segments: int. The number of types that the 'segment_ids' input can
+            take.
+        dtype: string 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"),
+        "segment_ids": np.array([[0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 0, 0]]),
+        "padding_mask": np.array([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0]]),
+    }
+
+    # Pretrained RoformerV2 encoder.
+    model = keras_hub.models.RoformerV2Backbone.from_preset("roformer_v2_base")
+    model(input_data)
+
+    # Randomly initialized RoformerV2 encoder with a custom config.
+    model = keras_hub.models.RoformerV2Backbone(
+        vocabulary_size=30552,
+        num_layers=4,
+        num_heads=4,
+        hidden_dim=256,
+        intermediate_dim=512,
+        head_size = 64,
+    )
+    model(input_data)
+    ```
+    """
+
+    def __init__(
+        self,
+        vocabulary_size,
+        num_layers,
+        num_heads,
+        hidden_dim,
+        intermediate_dim,
+        head_size,
+        use_bias=False,
+        activation="relu",
+        dropout=0.1,
+        num_segments=2,
+        dtype=None,
+        max_wavelength=10000,
+        **kwargs,
+    ):
+        # === Layers ===
+        self.token_embedding = ReversibleEmbedding(
+            input_dim=vocabulary_size,
+            output_dim=hidden_dim,
+            embeddings_initializer=roformer_kernel_initializer(),
+            dtype=dtype,
+            name="token_embedding",
+        )
+        self.segment_embedding = keras.layers.Embedding(
+            input_dim=num_segments,
+            output_dim=hidden_dim,
+            embeddings_initializer=roformer_kernel_initializer(),
+            dtype=dtype,
+            name="segment_embedding",
+        )
+        self.embeddings_add = keras.layers.Add(
+            dtype=dtype,
+            name="embeddings_add",
+        )
+        self.embeddings_layer_norm = RoformerNorm(
+            epsilon=keras.backend.epsilon(),
+            dtype=dtype,
+            name="embeddings_layer_norm",
+        )
+        self.embeddings_dropout = keras.layers.Dropout(
+            dropout,
+            dtype=dtype,
+            name="embeddings_dropout",
+        )
+        self.transformer_layers = []
+        for i in range(num_layers):
+            layer = RoformerV2Encoder(
+                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=roformer_kernel_initializer(),
+                dtype=dtype,
+                name=f"transformer_layer_{i}",
+            )
+            self.transformer_layers.append(layer)
+
+        # === Functional Model ===
+        token_id_input = keras.Input(
+            shape=(None,), dtype="int32", name="token_ids"
+        )
+        segment_id_input = keras.Input(
+            shape=(None,), dtype="int32", name="segment_ids"
+        )
+        attention_mask = keras.ops.not_equal(token_id_input, 0)
+        # Embed tokens, positions, and segment ids.
+        tokens = self.token_embedding(token_id_input)
+        segments = self.segment_embedding(segment_id_input)
+        # Sum, normalize and apply dropout to embeddings.
+        x = self.embeddings_add((tokens, segments))
+        x = self.embeddings_layer_norm(x)
+        x = self.embeddings_dropout(x)
+        for transformer_layer in self.transformer_layers:
+            x = transformer_layer(x, attention_mask=attention_mask)
+
+        super().__init__(
+            inputs={
+                "token_ids": token_id_input,
+                "segment_ids": segment_id_input,
+            },
+            outputs=x,
+            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.num_segments = num_segments
+        self.max_wavelength = max_wavelength
+        self.head_size = head_size
+        self.dropout = dropout
+        self.activation = activations.get(activation)
+        self.use_bias = use_bias
+        self.start_token_index = 0
+
+    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,
+                "num_segments": self.num_segments,
+                "max_wavelength": self.max_wavelength,
+                "head_size": self.head_size,
+                "use_bias": self.use_bias,
+                "activation": activations.serialize(self.activation),
+            }
+        )
+        return config