keras-hub 0.25.1__tar.gz → 0.26.0.dev0__tar.gz

{keras_hub-0.25.1 → keras_hub-0.26.0.dev0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: keras-hub
-Version: 0.25.1
+Version: 0.26.0.dev0
 Summary: Pretrained models for Keras.
 Author-email: Keras team <keras-users@googlegroups.com>
 License-Expression: Apache-2.0
@@ -8,7 +8,6 @@ Project-URL: Home, https://keras.io/keras_hub/
 Project-URL: Repository, https://github.com/keras-team/keras/keras_hub
 Classifier: Development Status :: 3 - Alpha
 Classifier: Programming Language :: Python :: 3
-Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3 :: Only
@@ -18,9 +17,9 @@ Classifier: Operating System :: MacOS
 Classifier: Intended Audience :: Science/Research
 Classifier: Topic :: Scientific/Engineering
 Classifier: Topic :: Software Development
-Requires-Python: >=3.10
+Requires-Python: >=3.11
 Description-Content-Type: text/markdown
-Requires-Dist: keras>=3.8
+Requires-Dist: keras>=3.13
 Requires-Dist: absl-py
 Requires-Dist: numpy
 Requires-Dist: packaging
@@ -31,7 +30,8 @@ Requires-Dist: tensorflow-text; platform_system != "Windows"
 
 # KerasHub: Multi-framework Pretrained Models
 [![](https://github.com/keras-team/keras-hub/workflows/Tests/badge.svg?branch=master)](https://github.com/keras-team/keras-hub/actions?query=workflow%3ATests+branch%3Amaster)
-![Python](https://img.shields.io/badge/python-v3.10.0+-success.svg)
+![Python](https://img.shields.io/badge/python-v3.11.0+-success.svg)
+[![Kaggle Models](https://img.shields.io/badge/Kaggle-Models-brightgreen?colorA=0099ff)](https://www.kaggle.com/organizations/keras/models)
 [![contributions welcome](https://img.shields.io/badge/contributions-welcome-brightgreen.svg?style=flat)](https://github.com/keras-team/keras-hub/issues)
 
 > [!IMPORTANT]
@@ -41,7 +41,7 @@ Requires-Dist: tensorflow-text; platform_system != "Windows"
 **KerasHub** is a pretrained modeling library that aims to be simple, flexible,
 and fast. The library provides [Keras 3](https://keras.io/keras_3/)
 implementations of popular model architectures, paired with a collection of
-pretrained checkpoints available on [Kaggle Models](https://kaggle.com/models/).
+pretrained checkpoints available on [Kaggle Models](https://www.kaggle.com/organizations/keras/models).
 Models can be used with text, image, and audio data for generation, classification,
 and many other built in tasks.
 
@@ -133,13 +133,6 @@ To install the latest KerasHub release with Keras 3, simply run:
 pip install --upgrade keras-hub
 ```
 
-Our text tokenizers are based on TensorFlow Text. Hence, if you are using any
-model which has language as a modality, you will have to run:
-
-```
-pip install --upgrade keras-hub[nlp]
-```
-
 To install the latest nightly changes for both KerasHub and Keras, you can use
 our nightly package.
 

{keras_hub-0.25.1 → keras_hub-0.26.0.dev0}/README.md

@@ -1,6 +1,7 @@
 # KerasHub: Multi-framework Pretrained Models
 [![](https://github.com/keras-team/keras-hub/workflows/Tests/badge.svg?branch=master)](https://github.com/keras-team/keras-hub/actions?query=workflow%3ATests+branch%3Amaster)
-![Python](https://img.shields.io/badge/python-v3.10.0+-success.svg)
+![Python](https://img.shields.io/badge/python-v3.11.0+-success.svg)
+[![Kaggle Models](https://img.shields.io/badge/Kaggle-Models-brightgreen?colorA=0099ff)](https://www.kaggle.com/organizations/keras/models)
 [![contributions welcome](https://img.shields.io/badge/contributions-welcome-brightgreen.svg?style=flat)](https://github.com/keras-team/keras-hub/issues)
 
 > [!IMPORTANT]
@@ -10,7 +11,7 @@
 **KerasHub** is a pretrained modeling library that aims to be simple, flexible,
 and fast. The library provides [Keras 3](https://keras.io/keras_3/)
 implementations of popular model architectures, paired with a collection of
-pretrained checkpoints available on [Kaggle Models](https://kaggle.com/models/).
+pretrained checkpoints available on [Kaggle Models](https://www.kaggle.com/organizations/keras/models).
 Models can be used with text, image, and audio data for generation, classification,
 and many other built in tasks.
 
@@ -102,13 +103,6 @@ To install the latest KerasHub release with Keras 3, simply run:
 pip install --upgrade keras-hub
 ```
 
-Our text tokenizers are based on TensorFlow Text. Hence, if you are using any
-model which has language as a modality, you will have to run:
-
-```
-pip install --upgrade keras-hub[nlp]
-```
-
 To install the latest nightly changes for both KerasHub and Keras, you can use
 our nightly package.
 
@@ -186,4 +180,4 @@ Thank you to all of our wonderful contributors!
 
 <a href="https://github.com/keras-team/keras-hub/graphs/contributors">
   <img src="https://contrib.rocks/image?repo=keras-team/keras-hub" />
-</a>
+</a>

{keras_hub-0.25.1 → keras_hub-0.26.0.dev0}/keras_hub/api/layers/__init__.py

@@ -138,6 +138,27 @@ from keras_hub.src.models.sam.sam_mask_decoder import (
 from keras_hub.src.models.sam.sam_prompt_encoder import (
     SAMPromptEncoder as SAMPromptEncoder,
 )
+from keras_hub.src.models.sam3.sam3_detr_decoder import (
+    SAM3DetrDecoder as SAM3DetrDecoder,
+)
+from keras_hub.src.models.sam3.sam3_detr_encoder import (
+    SAM3DetrEncoder as SAM3DetrEncoder,
+)
+from keras_hub.src.models.sam3.sam3_geometry_encoder import (
+    SAM3GeometryEncoder as SAM3GeometryEncoder,
+)
+from keras_hub.src.models.sam3.sam3_image_converter import (
+    SAM3ImageConverter as SAM3ImageConverter,
+)
+from keras_hub.src.models.sam3.sam3_mask_decoder import (
+    SAM3MaskDecoder as SAM3MaskDecoder,
+)
+from keras_hub.src.models.sam3.sam3_text_encoder import (
+    SAM3TextEncoder as SAM3TextEncoder,
+)
+from keras_hub.src.models.sam3.sam3_vision_encoder import (
+    SAM3VisionEncoder as SAM3VisionEncoder,
+)
 from keras_hub.src.models.segformer.segformer_image_converter import (
     SegFormerImageConverter as SegFormerImageConverter,
 )

{keras_hub-0.25.1 → keras_hub-0.26.0.dev0}/keras_hub/api/models/__init__.py

@@ -211,6 +211,12 @@ from keras_hub.src.models.distil_bert.distil_bert_text_classifier_preprocessor i
 from keras_hub.src.models.distil_bert.distil_bert_tokenizer import (
     DistilBertTokenizer as DistilBertTokenizer,
 )
+from keras_hub.src.models.edrec.edrec_backbone import (
+    EdRecBackbone as EdRecBackbone,
+)
+from keras_hub.src.models.edrec.edrec_seq2seq_lm import (
+    EdRecSeq2SeqLM as EdRecSeq2SeqLM,
+)
 from keras_hub.src.models.efficientnet.efficientnet_backbone import (
     EfficientNetBackbone as EfficientNetBackbone,
 )
@@ -629,6 +635,15 @@ from keras_hub.src.models.roformer_v2.roformer_v2_text_classifier_preprocessor i
 from keras_hub.src.models.roformer_v2.roformer_v2_tokenizer import (
     RoformerV2Tokenizer as RoformerV2Tokenizer,
 )
+from keras_hub.src.models.rwkv7.rwkv7_backbone import (
+    RWKV7Backbone as RWKV7Backbone,
+)
+from keras_hub.src.models.rwkv7.rwkv7_causal_lm import (
+    RWKV7CausalLM as RWKV7CausalLM,
+)
+from keras_hub.src.models.rwkv7.rwkv7_causal_lm_preprocessor import (
+    RWKV7CausalLMPreprocessor as RWKV7CausalLMPreprocessor,
+)
 from keras_hub.src.models.sam.sam_backbone import SAMBackbone as SAMBackbone
 from keras_hub.src.models.sam.sam_image_segmenter import (
     SAMImageSegmenter as SAMImageSegmenter,
@@ -636,6 +651,18 @@ from keras_hub.src.models.sam.sam_image_segmenter import (
 from keras_hub.src.models.sam.sam_image_segmenter_preprocessor import (
     SAMImageSegmenterPreprocessor as SAMImageSegmenterPreprocessor,
 )
+from keras_hub.src.models.sam3.sam3_pc_backbone import (
+    SAM3PromptableConceptBackbone as SAM3PromptableConceptBackbone,
+)
+from keras_hub.src.models.sam3.sam3_pc_image_segmenter import (
+    SAM3PromptableConceptImageSegmenter as SAM3PromptableConceptImageSegmenter,
+)
+from keras_hub.src.models.sam3.sam3_pc_image_segmenter_preprocessor import (
+    SAM3PromptableConceptImageSegmenterPreprocessor as SAM3PromptableConceptImageSegmenterPreprocessor,
+)
+from keras_hub.src.models.sam3.sam3_tokenizer import (
+    SAM3Tokenizer as SAM3Tokenizer,
+)
 from keras_hub.src.models.segformer.segformer_backbone import (
     SegFormerBackbone as SegFormerBackbone,
 )

{keras_hub-0.25.1 → keras_hub-0.26.0.dev0}/keras_hub/api/tokenizers/__init__.py

@@ -93,6 +93,12 @@ from keras_hub.src.models.roberta.roberta_tokenizer import (
 from keras_hub.src.models.roformer_v2.roformer_v2_tokenizer import (
     RoformerV2Tokenizer as RoformerV2Tokenizer,
 )
+from keras_hub.src.models.rwkv7.rwkv7_tokenizer import (
+    RWKVTokenizer as RWKVTokenizer,
+)
+from keras_hub.src.models.sam3.sam3_tokenizer import (
+    SAM3Tokenizer as SAM3Tokenizer,
+)
 from keras_hub.src.models.siglip.siglip_tokenizer import (
     SigLIPTokenizer as SigLIPTokenizer,
 )

{keras_hub-0.25.1 → keras_hub-0.26.0.dev0}/keras_hub/src/layers/modeling/non_max_supression.py

@@ -290,16 +290,19 @@ def non_max_suppression(
         "int32",
     )
     idx = ops.minimum(idx, num_boxes - 1)
+    idx = ops.cast(idx, "int32")
 
     index_offsets = ops.cast(ops.arange(batch_size) * num_boxes, "int32")
     take_along_axis_idx = ops.reshape(
         idx + ops.expand_dims(index_offsets, 1), [-1]
     )
+    take_along_axis_idx = ops.cast(take_along_axis_idx, "int32")
 
     if keras.backend.backend() != "tensorflow":
-        idx = ops.take_along_axis(
-            ops.reshape(sorted_indices, [-1]), take_along_axis_idx
+        sorted_indices_int = ops.cast(
+            ops.reshape(sorted_indices, [-1]), "int32"
         )
+        idx = ops.take_along_axis(sorted_indices_int, take_along_axis_idx)
     else:
         import tensorflow as tf
 

keras_hub-0.26.0.dev0/keras_hub/src/layers/modeling/reversible_embedding.py

@@ -0,0 +1,8 @@
+import keras
+
+from keras_hub.src.api_export import keras_hub_export
+
+
+@keras_hub_export("keras_hub.layers.ReversibleEmbedding")
+class ReversibleEmbedding(keras.layers.ReversibleEmbedding):
+    pass

{keras_hub-0.25.1 → keras_hub-0.26.0.dev0}/keras_hub/src/layers/modeling/token_and_position_embedding.py

@@ -1,10 +1,10 @@
 import keras
+from keras.layers import ReversibleEmbedding
+from keras.src.backend import get_keras_mask
+from keras.src.backend import set_keras_mask
 
 from keras_hub.src.api_export import keras_hub_export
 from keras_hub.src.layers.modeling.position_embedding import PositionEmbedding
-from keras_hub.src.layers.modeling.reversible_embedding import (
-    ReversibleEmbedding,
-)
 from keras_hub.src.utils.keras_utils import clone_initializer
 
 
@@ -128,10 +128,10 @@ class TokenAndPositionEmbedding(keras.layers.Layer):
             positions=positions,
         )
         outputs = embedded_tokens + embedded_positions
+        mask = get_keras_mask(embedded_tokens)
+        if mask is not None:
+            set_keras_mask(outputs, mask)
         return outputs
 
-    def compute_mask(self, inputs, mask=None):
-        return self.token_embedding.compute_mask(inputs, mask=mask)
-
     def compute_output_shape(self, input_shape):
         return tuple(input_shape) + (self.embedding_dim,)

{keras_hub-0.25.1 → keras_hub-0.26.0.dev0}/keras_hub/src/layers/modeling/transformer_layer_utils.py

@@ -1,11 +1,12 @@
 from absl import logging
 from keras import ops
+from keras.src.backend import get_keras_mask
 
 
 def _check_masks_shapes(inputs, padding_mask, attention_mask):
     mask = padding_mask
-    if hasattr(inputs, "_keras_mask") and mask is None:
-        mask = inputs._keras_mask
+    if mask is None:
+        mask = get_keras_mask(inputs)
     if mask is not None:
         if len(mask.shape) != 2:
             raise ValueError(
@@ -68,17 +69,16 @@ def merge_padding_and_attention_mask(
         returned mask is padding_mask with one additional axis.
     """
     _check_masks_shapes(inputs, padding_mask, attention_mask)
-    mask = padding_mask
-    if hasattr(inputs, "_keras_mask"):
-        if mask is None:
-            # If no padding mask is explicitly provided, we look for padding
-            # mask from the input data.
-            mask = inputs._keras_mask
-        else:
+    # We look for a padding mask from the input data.
+    mask = get_keras_mask(inputs)
+    # But if padding mask is explicitly provided, we use it.
+    if padding_mask is not None:
+        if mask is not None:
             logging.warning(
                 "You are explicitly setting `padding_mask` while the `inputs` "
                 "have built-in mask, so the built-in mask is ignored."
             )
+        mask = padding_mask
     if mask is not None:
         # Add an axis for broadcasting, the attention mask should be 2D
         # (not including the batch axis).

{keras_hub-0.25.1 → keras_hub-0.26.0.dev0}/keras_hub/src/layers/preprocessing/masked_lm_mask_generator.py

@@ -7,9 +7,11 @@ from keras_hub.src.utils.tensor_utils import preprocessing_function
 
 try:
     import tensorflow as tf
-    import tensorflow_text as tf_text
 except ImportError:
     tf = None
+try:
+    import tensorflow_text as tf_text
+except ImportError:
     tf_text = None
 
 

{keras_hub-0.25.1 → keras_hub-0.26.0.dev0}/keras_hub/src/layers/preprocessing/multi_segment_packer.py

@@ -8,9 +8,11 @@ from keras_hub.src.utils.tensor_utils import preprocessing_function
 
 try:
     import tensorflow as tf
-    import tensorflow_text as tf_text
 except ImportError:
     tf = None
+try:
+    import tensorflow_text as tf_text
+except ImportError:
     tf_text = None
 
 

{keras_hub-0.25.1 → keras_hub-0.26.0.dev0}/keras_hub/src/models/albert/albert_backbone.py

@@ -1,10 +1,8 @@
 import keras
+from keras.layers import ReversibleEmbedding
 
 from keras_hub.src.api_export import keras_hub_export
 from keras_hub.src.layers.modeling.position_embedding import PositionEmbedding
-from keras_hub.src.layers.modeling.reversible_embedding import (
-    ReversibleEmbedding,
-)
 from keras_hub.src.layers.modeling.transformer_encoder import TransformerEncoder
 from keras_hub.src.models.backbone import Backbone
 from keras_hub.src.utils.keras_utils import gelu_approximate

{keras_hub-0.25.1 → keras_hub-0.26.0.dev0}/keras_hub/src/models/backbone.py

@@ -107,6 +107,9 @@ class Backbone(keras.Model):
     def from_config(cls, config):
         # The default `from_config()` for functional models will return a
         # vanilla `keras.Model`. We override it to get a subclass instance back.
+        config = config.copy()
+        if "dtype" in config and isinstance(config["dtype"], dict):
+            config["dtype"] = keras.dtype_policies.get(config["dtype"])
         return cls(**config)
 
     @classproperty

{keras_hub-0.25.1 → keras_hub-0.26.0.dev0}/keras_hub/src/models/bart/bart_backbone.py

@@ -1,10 +1,8 @@
 import keras
+from keras.layers import ReversibleEmbedding
 
 from keras_hub.src.api_export import keras_hub_export
 from keras_hub.src.layers.modeling.position_embedding import PositionEmbedding
-from keras_hub.src.layers.modeling.reversible_embedding import (
-    ReversibleEmbedding,
-)
 from keras_hub.src.layers.modeling.transformer_decoder import TransformerDecoder
 from keras_hub.src.layers.modeling.transformer_encoder import TransformerEncoder
 from keras_hub.src.models.backbone import Backbone

{keras_hub-0.25.1 → keras_hub-0.26.0.dev0}/keras_hub/src/models/bert/bert_backbone.py

@@ -1,10 +1,8 @@
 import keras
+from keras.layers import ReversibleEmbedding
 
 from keras_hub.src.api_export import keras_hub_export
 from keras_hub.src.layers.modeling.position_embedding import PositionEmbedding
-from keras_hub.src.layers.modeling.reversible_embedding import (
-    ReversibleEmbedding,
-)
 from keras_hub.src.layers.modeling.transformer_encoder import TransformerEncoder
 from keras_hub.src.models.backbone import Backbone
 from keras_hub.src.utils.keras_utils import gelu_approximate
@@ -35,7 +33,7 @@ class BertBackbone(Backbone):
         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.
+            The hidden_dim 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.

{keras_hub-0.25.1 → keras_hub-0.26.0.dev0}/keras_hub/src/models/bloom/bloom_backbone.py

@@ -1,9 +1,7 @@
 import keras
+from keras.layers import ReversibleEmbedding
 
 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.bloom.bloom_decoder import BloomDecoder
 

{keras_hub-0.25.1 → keras_hub-0.26.0.dev0}/keras_hub/src/models/causal_lm.py

@@ -196,7 +196,7 @@ class CausalLM(Task):
 
                 # Create an explicit tuple of all variable state.
                 state = (
-                    self.sampler.variables,
+                    [v.value for v in self.sampler.variables],
                     # Use the explicit variable.value to preserve the
                     # sharding spec of distribution.
                     [v.value for v in self.trainable_variables],
@@ -431,7 +431,7 @@ class CausalLM(Task):
         self.generate_function = None
 
     def get_quantization_layer_structure(self, mode):
-        if mode != "gptq":
+        if mode not in ["gptq", "awq"]:
             return None
 
         backbone = self.backbone

{keras_hub-0.25.1 → keras_hub-0.26.0.dev0}/keras_hub/src/models/deberta_v3/deberta_v3_backbone.py

@@ -1,9 +1,7 @@
 import keras
+from keras.layers import ReversibleEmbedding
 
 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.deberta_v3.disentangled_attention_encoder import (
     DisentangledAttentionEncoder,

keras_hub-0.26.0.dev0/keras_hub/src/models/edrec/edrec_backbone.py

@@ -0,0 +1,147 @@
+import keras
+
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.models.backbone import Backbone
+from keras_hub.src.models.edrec.edrec_layers import EdRecDecoderBlock
+from keras_hub.src.models.edrec.edrec_layers import EdRecEncoderBlock
+
+
+@keras_hub_export("keras_hub.models.EdRecBackbone")
+class EdRecBackbone(Backbone):
+    """EdRec Backbone model.
+
+    Args:
+        vocab_size: int, size of the vocabulary.
+        num_layers_enc: int, number of encoder layers.
+        num_layers_dec: int, number of decoder layers.
+        hidden_dim: int, hidden dimension (d_model).
+        intermediate_dim: int, intermediate dimension (d_ff).
+        num_heads: int, number of attention heads.
+        dropout: float, dropout rate.
+        epsilon: float, epsilon for simple RMSNorm.
+    """
+
+    def __init__(
+        self,
+        vocab_size,
+        num_layers_enc,
+        num_layers_dec,
+        hidden_dim,
+        intermediate_dim,
+        num_heads,
+        dropout=0.0,
+        epsilon=1e-6,
+        dtype=None,
+        **kwargs,
+    ):
+        # === Layers ===
+        self.embedding = keras.layers.Embedding(
+            input_dim=vocab_size,
+            output_dim=hidden_dim,
+            dtype=dtype,
+            name="embedding",
+        )
+        self.encoder_layers = []
+        for i in range(num_layers_enc):
+            self.encoder_layers.append(
+                EdRecEncoderBlock(
+                    hidden_dim=hidden_dim,
+                    num_heads=num_heads,
+                    intermediate_dim=intermediate_dim,
+                    dropout_rate=dropout,
+                    epsilon=epsilon,
+                    dtype=dtype,
+                    name=f"encoder_layer_{i}",
+                )
+            )
+        self.decoder_layers = []
+        for i in range(num_layers_dec):
+            self.decoder_layers.append(
+                EdRecDecoderBlock(
+                    hidden_dim=hidden_dim,
+                    num_heads=num_heads,
+                    intermediate_dim=intermediate_dim,
+                    dropout_rate=dropout,
+                    epsilon=epsilon,
+                    dtype=dtype,
+                    name=f"decoder_layer_{i}",
+                )
+            )
+
+        # === Functional Model ===
+        encoder_token_ids = keras.Input(
+            shape=(None,), dtype="int32", name="encoder_token_ids"
+        )
+        decoder_token_ids = keras.Input(
+            shape=(None,), dtype="int32", name="decoder_token_ids"
+        )
+        encoder_padding_mask = keras.Input(
+            shape=(None,), dtype="bool", name="encoder_padding_mask"
+        )
+        decoder_padding_mask = keras.Input(
+            shape=(None,), dtype="bool", name="decoder_padding_mask"
+        )
+
+        # Encoder
+        x_enc = self.embedding(encoder_token_ids)
+
+        for layer in self.encoder_layers:
+            x_enc = layer(
+                x_enc,
+                padding_mask=encoder_padding_mask,
+            )
+
+        # Decoder
+        x_dec = self.embedding(decoder_token_ids)
+        for layer in self.decoder_layers:
+            x_dec, _, _ = layer(
+                x_dec,
+                encoder_outputs=x_enc,
+                decoder_padding_mask=decoder_padding_mask,
+                encoder_padding_mask=encoder_padding_mask,
+            )
+
+        super().__init__(
+            inputs={
+                "encoder_token_ids": encoder_token_ids,
+                "decoder_token_ids": decoder_token_ids,
+                "encoder_padding_mask": encoder_padding_mask,
+                "decoder_padding_mask": decoder_padding_mask,
+            },
+            outputs={
+                "encoder_sequence_output": x_enc,
+                "decoder_sequence_output": x_dec,
+            },
+            dtype=dtype,
+            **kwargs,
+        )
+
+        # === Config ===
+        self.vocab_size = vocab_size
+        self.num_layers_enc = num_layers_enc
+        self.num_layers_dec = num_layers_dec
+        self.hidden_dim = hidden_dim
+        self.intermediate_dim = intermediate_dim
+        self.num_heads = num_heads
+        self.dropout = dropout
+        self.epsilon = epsilon
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "vocab_size": self.vocab_size,
+                "num_layers_enc": self.num_layers_enc,
+                "num_layers_dec": self.num_layers_dec,
+                "hidden_dim": self.hidden_dim,
+                "intermediate_dim": self.intermediate_dim,
+                "num_heads": self.num_heads,
+                "dropout": self.dropout,
+                "epsilon": self.epsilon,
+            }
+        )
+        return config
+
+    @property
+    def token_embedding(self):
+        return self.embedding