keras-hub-nightly 0.15.0.dev20240823171555__py3-none-any.whl → 0.16.0.dev2024092017__py3-none-any.whl

keras_hub/src/utils/transformers/convert_bart.py

@@ -0,0 +1,373 @@
+# Copyright 2024 The KerasHub Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import numpy as np
+
+from keras_hub.src.models.bart.bart_backbone import BartBackbone
+from keras_hub.src.utils.preset_utils import get_file
+
+backbone_cls = BartBackbone
+
+
+def convert_backbone_config(transformers_config):
+    return {
+        "vocabulary_size": transformers_config["vocab_size"],
+        "num_layers": transformers_config["num_hidden_layers"],
+        "num_heads": transformers_config["encoder_attention_heads"],
+        "hidden_dim": transformers_config["d_model"],
+        "intermediate_dim": transformers_config["encoder_ffn_dim"],
+        "dropout": transformers_config["dropout"],
+        "max_sequence_length": transformers_config["max_position_embeddings"],
+    }
+
+
+def convert_weights(backbone, loader, transformers_config):
+    # Embeddings
+    loader.port_weight(
+        keras_variable=backbone.token_embedding.embeddings,
+        hf_weight_key="shared.weight",
+    )
+    loader.port_weight(
+        keras_variable=backbone.encoder_position_embedding.position_embeddings,
+        hf_weight_key="encoder.embed_positions.weight",
+        hook_fn=lambda hf_tensor, keras_shape: np.reshape(
+            hf_tensor[2:, :], keras_shape
+        ),
+    )
+    loader.port_weight(
+        keras_variable=backbone.decoder_position_embedding.position_embeddings,
+        hf_weight_key="decoder.embed_positions.weight",
+        hook_fn=lambda hf_tensor, keras_shape: np.reshape(
+            hf_tensor[2:, :], keras_shape
+        ),
+    )
+
+    # Encoder blocks
+    for index in range(backbone.num_layers):
+        encoder_layer = backbone.encoder_transformer_layers[index]
+        encoder_self_attention = encoder_layer._self_attention_layer
+        hf_encoder_prefix = f"encoder.layers.{index}"
+
+        # Norm layers
+        loader.port_weight(
+            keras_variable=encoder_layer._self_attention_layer_norm.gamma,
+            hf_weight_key=f"{hf_encoder_prefix}.self_attn_layer_norm.weight",
+        )
+        loader.port_weight(
+            keras_variable=encoder_layer._self_attention_layer_norm.beta,
+            hf_weight_key=f"{hf_encoder_prefix}.self_attn_layer_norm.bias",
+        )
+        loader.port_weight(
+            keras_variable=encoder_layer._feedforward_layer_norm.gamma,
+            hf_weight_key=f"{hf_encoder_prefix}.final_layer_norm.weight",
+        )
+        loader.port_weight(
+            keras_variable=encoder_layer._feedforward_layer_norm.beta,
+            hf_weight_key=f"{hf_encoder_prefix}.final_layer_norm.bias",
+        )
+
+        # Self Attention layers
+        # Query
+        loader.port_weight(
+            keras_variable=encoder_self_attention.query_dense.kernel,
+            hf_weight_key=f"{hf_encoder_prefix}.self_attn.q_proj.weight",
+            hook_fn=lambda hf_tensor, keras_shape: np.reshape(
+                np.transpose(hf_tensor), keras_shape
+            ),
+        )
+        loader.port_weight(
+            keras_variable=encoder_self_attention.query_dense.bias,
+            hf_weight_key=f"{hf_encoder_prefix}.self_attn.q_proj.bias",
+            hook_fn=lambda hf_tensor, keras_shape: np.reshape(
+                np.transpose(hf_tensor), keras_shape
+            ),
+        )
+
+        # Key
+        loader.port_weight(
+            keras_variable=encoder_self_attention.key_dense.kernel,
+            hf_weight_key=f"{hf_encoder_prefix}.self_attn.k_proj.weight",
+            hook_fn=lambda hf_tensor, keras_shape: np.reshape(
+                np.transpose(hf_tensor), keras_shape
+            ),
+        )
+        loader.port_weight(
+            keras_variable=encoder_self_attention.key_dense.bias,
+            hf_weight_key=f"{hf_encoder_prefix}.self_attn.k_proj.bias",
+            hook_fn=lambda hf_tensor, keras_shape: np.reshape(
+                np.transpose(hf_tensor), keras_shape
+            ),
+        )
+
+        # Value
+        loader.port_weight(
+            keras_variable=encoder_self_attention.value_dense.kernel,
+            hf_weight_key=f"{hf_encoder_prefix}.self_attn.v_proj.weight",
+            hook_fn=lambda hf_tensor, keras_shape: np.reshape(
+                np.transpose(hf_tensor), keras_shape
+            ),
+        )
+        loader.port_weight(
+            keras_variable=encoder_self_attention.value_dense.bias,
+            hf_weight_key=f"{hf_encoder_prefix}.self_attn.v_proj.bias",
+            hook_fn=lambda hf_tensor, keras_shape: np.reshape(
+                np.transpose(hf_tensor), keras_shape
+            ),
+        )
+
+        # Output
+        loader.port_weight(
+            keras_variable=encoder_self_attention.output_dense.kernel,
+            hf_weight_key=f"{hf_encoder_prefix}.self_attn.out_proj.weight",
+            hook_fn=lambda hf_tensor, keras_shape: np.reshape(
+                np.transpose(hf_tensor), keras_shape
+            ),
+        )
+        loader.port_weight(
+            keras_variable=encoder_self_attention.output_dense.bias,
+            hf_weight_key=f"{hf_encoder_prefix}.self_attn.out_proj.bias",
+            hook_fn=lambda hf_tensor, keras_shape: np.reshape(
+                np.transpose(hf_tensor), keras_shape
+            ),
+        )
+
+        # MLP layers
+        loader.port_weight(
+            keras_variable=encoder_layer._feedforward_intermediate_dense.kernel,
+            hf_weight_key=f"{hf_encoder_prefix}.fc1.weight",
+            hook_fn=lambda hf_tensor, _: np.transpose(hf_tensor, axes=(1, 0)),
+        )
+        loader.port_weight(
+            keras_variable=encoder_layer._feedforward_intermediate_dense.bias,
+            hf_weight_key=f"{hf_encoder_prefix}.fc1.bias",
+        )
+        loader.port_weight(
+            keras_variable=encoder_layer._feedforward_output_dense.kernel,
+            hf_weight_key=f"{hf_encoder_prefix}.fc2.weight",
+            hook_fn=lambda hf_tensor, _: np.transpose(hf_tensor, axes=(1, 0)),
+        )
+        loader.port_weight(
+            keras_variable=encoder_layer._feedforward_output_dense.bias,
+            hf_weight_key=f"{hf_encoder_prefix}.fc2.bias",
+        )
+
+    # Decoder blocks
+    for index in range(backbone.num_layers):
+        decoder_layer = backbone.decoder_transformer_layers[index]
+        decoder_self_attention = decoder_layer._self_attention_layer
+        decoder_cross_attention = decoder_layer._cross_attention_layer
+        hf_decoder_prefix = f"decoder.layers.{index}"
+
+        # Norm layers
+        loader.port_weight(
+            keras_variable=decoder_layer._self_attention_layer_norm.gamma,
+            hf_weight_key=f"{hf_decoder_prefix}.self_attn_layer_norm.weight",
+        )
+        loader.port_weight(
+            keras_variable=decoder_layer._self_attention_layer_norm.beta,
+            hf_weight_key=f"{hf_decoder_prefix}.self_attn_layer_norm.bias",
+        )
+        loader.port_weight(
+            keras_variable=decoder_layer._feedforward_layer_norm.gamma,
+            hf_weight_key=f"{hf_decoder_prefix}.final_layer_norm.weight",
+        )
+        loader.port_weight(
+            keras_variable=decoder_layer._feedforward_layer_norm.beta,
+            hf_weight_key=f"{hf_decoder_prefix}.final_layer_norm.bias",
+        )
+        loader.port_weight(
+            keras_variable=decoder_layer._cross_attention_layer_norm.gamma,
+            hf_weight_key=f"{hf_decoder_prefix}.encoder_attn_layer_norm.weight",
+        )
+        loader.port_weight(
+            keras_variable=decoder_layer._cross_attention_layer_norm.beta,
+            hf_weight_key=f"{hf_decoder_prefix}.encoder_attn_layer_norm.bias",
+        )
+
+        # Self Attention layers
+        # Query
+        loader.port_weight(
+            keras_variable=decoder_self_attention.query_dense.kernel,
+            hf_weight_key=f"{hf_decoder_prefix}.self_attn.q_proj.weight",
+            hook_fn=lambda hf_tensor, keras_shape: np.reshape(
+                np.transpose(hf_tensor), keras_shape
+            ),
+        )
+        loader.port_weight(
+            keras_variable=decoder_self_attention.query_dense.bias,
+            hf_weight_key=f"{hf_decoder_prefix}.self_attn.q_proj.bias",
+            hook_fn=lambda hf_tensor, keras_shape: np.reshape(
+                np.transpose(hf_tensor), keras_shape
+            ),
+        )
+
+        # Key
+        loader.port_weight(
+            keras_variable=decoder_self_attention.key_dense.kernel,
+            hf_weight_key=f"{hf_decoder_prefix}.self_attn.k_proj.weight",
+            hook_fn=lambda hf_tensor, keras_shape: np.reshape(
+                np.transpose(hf_tensor), keras_shape
+            ),
+        )
+        loader.port_weight(
+            keras_variable=decoder_self_attention.key_dense.bias,
+            hf_weight_key=f"{hf_decoder_prefix}.self_attn.k_proj.bias",
+            hook_fn=lambda hf_tensor, keras_shape: np.reshape(
+                np.transpose(hf_tensor), keras_shape
+            ),
+        )
+
+        # Value
+        loader.port_weight(
+            keras_variable=decoder_self_attention.value_dense.kernel,
+            hf_weight_key=f"{hf_decoder_prefix}.self_attn.v_proj.weight",
+            hook_fn=lambda hf_tensor, keras_shape: np.reshape(
+                np.transpose(hf_tensor), keras_shape
+            ),
+        )
+        loader.port_weight(
+            keras_variable=decoder_self_attention.value_dense.bias,
+            hf_weight_key=f"{hf_decoder_prefix}.self_attn.v_proj.bias",
+            hook_fn=lambda hf_tensor, keras_shape: np.reshape(
+                np.transpose(hf_tensor), keras_shape
+            ),
+        )
+
+        # Output
+        loader.port_weight(
+            keras_variable=decoder_self_attention.output_dense.kernel,
+            hf_weight_key=f"{hf_decoder_prefix}.self_attn.out_proj.weight",
+            hook_fn=lambda hf_tensor, keras_shape: np.reshape(
+                np.transpose(hf_tensor), keras_shape
+            ),
+        )
+        loader.port_weight(
+            keras_variable=decoder_self_attention.output_dense.bias,
+            hf_weight_key=f"{hf_decoder_prefix}.self_attn.out_proj.bias",
+            hook_fn=lambda hf_tensor, keras_shape: np.reshape(
+                np.transpose(hf_tensor), keras_shape
+            ),
+        )
+
+        # MLP layers
+        loader.port_weight(
+            keras_variable=decoder_layer._feedforward_intermediate_dense.kernel,
+            hf_weight_key=f"{hf_decoder_prefix}.fc1.weight",
+            hook_fn=lambda hf_tensor, _: np.transpose(hf_tensor, axes=(1, 0)),
+        )
+        loader.port_weight(
+            keras_variable=decoder_layer._feedforward_intermediate_dense.bias,
+            hf_weight_key=f"{hf_decoder_prefix}.fc1.bias",
+        )
+        loader.port_weight(
+            keras_variable=decoder_layer._feedforward_output_dense.kernel,
+            hf_weight_key=f"{hf_decoder_prefix}.fc2.weight",
+            hook_fn=lambda hf_tensor, _: np.transpose(hf_tensor, axes=(1, 0)),
+        )
+        loader.port_weight(
+            keras_variable=decoder_layer._feedforward_output_dense.bias,
+            hf_weight_key=f"{hf_decoder_prefix}.fc2.bias",
+        )
+
+        # Cross Attention Layers
+        # Query
+        loader.port_weight(
+            keras_variable=decoder_cross_attention.query_dense.kernel,
+            hf_weight_key=f"{hf_decoder_prefix}.encoder_attn.q_proj.weight",
+            hook_fn=lambda hf_tensor, keras_shape: np.reshape(
+                np.transpose(hf_tensor), keras_shape
+            ),
+        )
+        loader.port_weight(
+            keras_variable=decoder_cross_attention.query_dense.bias,
+            hf_weight_key=f"{hf_decoder_prefix}.encoder_attn.q_proj.bias",
+            hook_fn=lambda hf_tensor, keras_shape: np.reshape(
+                np.transpose(hf_tensor), keras_shape
+            ),
+        )
+
+        # Key
+        loader.port_weight(
+            keras_variable=decoder_cross_attention.key_dense.kernel,
+            hf_weight_key=f"{hf_decoder_prefix}.encoder_attn.k_proj.weight",
+            hook_fn=lambda hf_tensor, keras_shape: np.reshape(
+                np.transpose(hf_tensor), keras_shape
+            ),
+        )
+        loader.port_weight(
+            keras_variable=decoder_cross_attention.key_dense.bias,
+            hf_weight_key=f"{hf_decoder_prefix}.encoder_attn.k_proj.bias",
+            hook_fn=lambda hf_tensor, keras_shape: np.reshape(
+                np.transpose(hf_tensor), keras_shape
+            ),
+        )
+
+        # Value
+        loader.port_weight(
+            keras_variable=decoder_cross_attention.value_dense.kernel,
+            hf_weight_key=f"{hf_decoder_prefix}.encoder_attn.v_proj.weight",
+            hook_fn=lambda hf_tensor, keras_shape: np.reshape(
+                np.transpose(hf_tensor), keras_shape
+            ),
+        )
+        loader.port_weight(
+            keras_variable=decoder_cross_attention.value_dense.bias,
+            hf_weight_key=f"{hf_decoder_prefix}.encoder_attn.v_proj.bias",
+            hook_fn=lambda hf_tensor, keras_shape: np.reshape(
+                np.transpose(hf_tensor), keras_shape
+            ),
+        )
+
+        # Output
+        loader.port_weight(
+            keras_variable=decoder_cross_attention.output_dense.kernel,
+            hf_weight_key=f"{hf_decoder_prefix}.encoder_attn.out_proj.weight",
+            hook_fn=lambda hf_tensor, keras_shape: np.reshape(
+                np.transpose(hf_tensor), keras_shape
+            ),
+        )
+        loader.port_weight(
+            keras_variable=decoder_cross_attention.output_dense.bias,
+            hf_weight_key=f"{hf_decoder_prefix}.encoder_attn.out_proj.bias",
+            hook_fn=lambda hf_tensor, keras_shape: np.reshape(
+                np.transpose(hf_tensor), keras_shape
+            ),
+        )
+
+    # Normalization
+    loader.port_weight(
+        keras_variable=backbone.encoder_embeddings_layer_norm.gamma,
+        hf_weight_key="encoder.layernorm_embedding.weight",
+    )
+    loader.port_weight(
+        keras_variable=backbone.encoder_embeddings_layer_norm.beta,
+        hf_weight_key="encoder.layernorm_embedding.bias",
+    )
+    loader.port_weight(
+        keras_variable=backbone.decoder_embeddings_layer_norm.gamma,
+        hf_weight_key="decoder.layernorm_embedding.weight",
+    )
+    loader.port_weight(
+        keras_variable=backbone.decoder_embeddings_layer_norm.beta,
+        hf_weight_key="decoder.layernorm_embedding.bias",
+    )
+
+
+def convert_tokenizer(cls, preset, **kwargs):
+    vocab_file = get_file(preset, "vocab.json")
+    merges_file = get_file(preset, "merges.txt")
+    return cls(
+        vocabulary=vocab_file,
+        merges=merges_file,
+        **kwargs,
+    )

keras_hub/src/utils/transformers/convert_bert.py

@@ -13,12 +13,12 @@
 # limitations under the License.
 import numpy as np
 
-from keras_hub.src.utils.preset_utils import HF_CONFIG_FILE
+from keras_hub.src.models.bert.bert_backbone import BertBackbone
 from keras_hub.src.utils.preset_utils import HF_TOKENIZER_CONFIG_FILE
 from keras_hub.src.utils.preset_utils import get_file
-from keras_hub.src.utils.preset_utils import jax_memory_cleanup
-from keras_hub.src.utils.preset_utils import load_config
-from keras_hub.src.utils.transformers.safetensor_utils import SafetensorLoader
+from keras_hub.src.utils.preset_utils import load_json
+
+backbone_cls = BertBackbone
 
 
 def convert_backbone_config(transformers_config):
@@ -154,20 +154,10 @@ def convert_weights(backbone, loader, transformers_config):
     )
 
 
-def load_bert_backbone(cls, preset, load_weights):
-    transformers_config = load_config(preset, HF_CONFIG_FILE)
-    keras_config = convert_backbone_config(transformers_config)
-    backbone = cls(**keras_config)
-    if load_weights:
-        jax_memory_cleanup(backbone)
-        with SafetensorLoader(preset) as loader:
-            convert_weights(backbone, loader, transformers_config)
-    return backbone
-
-
-def load_bert_tokenizer(cls, preset):
-    transformers_config = load_config(preset, HF_TOKENIZER_CONFIG_FILE)
+def convert_tokenizer(cls, preset, **kwargs):
+    transformers_config = load_json(preset, HF_TOKENIZER_CONFIG_FILE)
     return cls(
         get_file(preset, "vocab.txt"),
         lowercase=transformers_config["do_lower_case"],
+        **kwargs,
     )

keras_hub/src/utils/transformers/convert_distilbert.py

@@ -13,12 +13,14 @@
 # limitations under the License.
 import numpy as np
 
-from keras_hub.src.utils.preset_utils import HF_CONFIG_FILE
+from keras_hub.src.models.distil_bert.distil_bert_backbone import (
+    DistilBertBackbone,
+)
 from keras_hub.src.utils.preset_utils import HF_TOKENIZER_CONFIG_FILE
 from keras_hub.src.utils.preset_utils import get_file
-from keras_hub.src.utils.preset_utils import jax_memory_cleanup
-from keras_hub.src.utils.preset_utils import load_config
-from keras_hub.src.utils.transformers.safetensor_utils import SafetensorLoader
+from keras_hub.src.utils.preset_utils import load_json
+
+backbone_cls = DistilBertBackbone
 
 
 def convert_backbone_config(transformers_config):
@@ -33,7 +35,7 @@ def convert_backbone_config(transformers_config):
     }
 
 
-def convert_weights(backbone, loader):
+def convert_weights(backbone, loader, transformers_config):
     # Embeddings
     loader.port_weight(
         keras_variable=backbone.get_layer(
@@ -162,23 +164,11 @@ def convert_weights(backbone, loader):
         hf_weight_key="distilbert.embeddings.LayerNorm.bias",
     )
 
-    return backbone
-
-
-def load_distilbert_backbone(cls, preset, load_weights):
-    transformers_config = load_config(preset, HF_CONFIG_FILE)
-    keras_config = convert_backbone_config(transformers_config)
-    backbone = cls(**keras_config)
-    if load_weights:
-        jax_memory_cleanup(backbone)
-        with SafetensorLoader(preset) as loader:
-            convert_weights(backbone, loader)
-    return backbone
-
 
-def load_distilbert_tokenizer(cls, preset):
-    transformers_config = load_config(preset, HF_TOKENIZER_CONFIG_FILE)
+def convert_tokenizer(cls, preset, **kwargs):
+    transformers_config = load_json(preset, HF_TOKENIZER_CONFIG_FILE)
     return cls(
         get_file(preset, "vocab.txt"),
         lowercase=transformers_config["do_lower_case"],
+        **kwargs,
     )

keras_hub/src/utils/transformers/convert_gemma.py

@@ -13,11 +13,10 @@
 # limitations under the License.
 import numpy as np
 
-from keras_hub.src.utils.preset_utils import HF_CONFIG_FILE
+from keras_hub.src.models.gemma.gemma_backbone import GemmaBackbone
 from keras_hub.src.utils.preset_utils import get_file
-from keras_hub.src.utils.preset_utils import jax_memory_cleanup
-from keras_hub.src.utils.preset_utils import load_config
-from keras_hub.src.utils.transformers.safetensor_utils import SafetensorLoader
+
+backbone_cls = GemmaBackbone
 
 
 def convert_backbone_config(transformers_config):
@@ -169,19 +168,6 @@ def convert_weights(backbone, loader, transformers_config):
         hf_weight_key="model.norm.weight",
     )
 
-    return backbone
-
-
-def load_gemma_backbone(cls, preset, load_weights):
-    transformers_config = load_config(preset, HF_CONFIG_FILE)
-    keras_config = convert_backbone_config(transformers_config)
-    backbone = cls(**keras_config)
-    if load_weights:
-        jax_memory_cleanup(backbone)
-        with SafetensorLoader(preset) as loader:
-            convert_weights(backbone, loader, transformers_config)
-    return backbone
-
 
-def load_gemma_tokenizer(cls, preset):
-    return cls(get_file(preset, "tokenizer.model"))
+def convert_tokenizer(cls, preset, **kwargs):
+    return cls(get_file(preset, "tokenizer.model"), **kwargs)

keras_hub/src/utils/transformers/convert_gpt2.py

@@ -13,11 +13,10 @@
 # limitations under the License.
 import numpy as np
 
-from keras_hub.src.utils.preset_utils import HF_CONFIG_FILE
+from keras_hub.src.models.gpt2.gpt2_backbone import GPT2Backbone
 from keras_hub.src.utils.preset_utils import get_file
-from keras_hub.src.utils.preset_utils import jax_memory_cleanup
-from keras_hub.src.utils.preset_utils import load_config
-from keras_hub.src.utils.transformers.safetensor_utils import SafetensorLoader
+
+backbone_cls = GPT2Backbone
 
 
 def convert_backbone_config(transformers_config):
@@ -163,24 +162,12 @@ def convert_weights(backbone, loader, transformers_config):
         hf_weight_key="ln_f.bias",
     )
 
-    return backbone
-
-
-def load_gpt2_backbone(cls, preset, load_weights):
-    transformers_config = load_config(preset, HF_CONFIG_FILE)
-    keras_config = convert_backbone_config(transformers_config)
-    backbone = cls(**keras_config)
-    if load_weights:
-        jax_memory_cleanup(backbone)
-        with SafetensorLoader(preset) as loader:
-            convert_weights(backbone, loader, transformers_config)
-    return backbone
-
 
-def load_gpt2_tokenizer(cls, preset):
+def convert_tokenizer(cls, preset, **kwargs):
     vocab_file = get_file(preset, "vocab.json")
     merges_file = get_file(preset, "merges.txt")
     return cls(
         vocabulary=vocab_file,
         merges=merges_file,
+        **kwargs,
     )

keras_hub/src/utils/transformers/convert_llama3.py

@@ -13,10 +13,10 @@
 # limitations under the License.
 import numpy as np
 
-from keras_hub.src.utils.preset_utils import HF_CONFIG_FILE
-from keras_hub.src.utils.preset_utils import jax_memory_cleanup
-from keras_hub.src.utils.preset_utils import load_config
-from keras_hub.src.utils.transformers.safetensor_utils import SafetensorLoader
+from keras_hub.src.models.llama3.llama3_backbone import Llama3Backbone
+from keras_hub.src.utils.preset_utils import load_json
+
+backbone_cls = Llama3Backbone
 
 
 def convert_backbone_config(transformers_config):
@@ -111,19 +111,8 @@ def convert_weights(backbone, loader, transformers_config):
     return backbone
 
 
-def load_llama3_backbone(cls, preset, load_weights):
-    transformers_config = load_config(preset, HF_CONFIG_FILE)
-    keras_config = convert_backbone_config(transformers_config)
-    backbone = cls(**keras_config)
-    if load_weights:
-        jax_memory_cleanup(backbone)
-        with SafetensorLoader(preset) as loader:
-            convert_weights(backbone, loader, transformers_config)
-    return backbone
-
-
-def load_llama3_tokenizer(cls, preset):
-    tokenizer_config = load_config(preset, "tokenizer.json")
+def convert_tokenizer(cls, preset, **kwargs):
+    tokenizer_config = load_json(preset, "tokenizer.json")
     vocab = tokenizer_config["model"]["vocab"]
     merges = tokenizer_config["model"]["merges"]
 
@@ -133,4 +122,4 @@ def load_llama3_tokenizer(cls, preset):
     vocab[bot["content"]] = bot["id"]
     vocab[eot["content"]] = eot["id"]
 
-    return cls(vocabulary=vocab, merges=merges)
+    return cls(vocabulary=vocab, merges=merges, **kwargs)