Add SafeTensors Conversion Support for Qwen2/2.5

LakshmiKalaKadali · LakshmiKalaKadali · commit 502488e91cc1 · 2025-12-12T20:40:31.000+05:30
Models
diff --git a/keras_hub/src/utils/transformers/export/hf_exporter.py b/keras_hub/src/utils/transformers/export/hf_exporter.py
@@ -5,34 +5,43 @@
 
 import keras
 
+# --- Gemma Utils ---
 from keras_hub.src.utils.transformers.export.gemma import get_gemma_config
 from keras_hub.src.utils.transformers.export.gemma import (
     get_gemma_tokenizer_config,
 )
 from keras_hub.src.utils.transformers.export.gemma import get_gemma_weights_map
+
+# --- GPT-2 Utils ---
 from keras_hub.src.utils.transformers.export.gpt2 import get_gpt2_config
 from keras_hub.src.utils.transformers.export.gpt2 import (
     get_gpt2_tokenizer_config,
 )
 from keras_hub.src.utils.transformers.export.gpt2 import get_gpt2_weights_map
 
+# --- Qwen Utils ---
+from keras_hub.src.utils.transformers.export.qwen import get_qwen_config
+from keras_hub.src.utils.transformers.export.qwen import (
+    get_qwen_tokenizer_config,
+)
+from keras_hub.src.utils.transformers.export.qwen import get_qwen_weights_map
+
 MODEL_CONFIGS = {
     "GemmaBackbone": get_gemma_config,
     "GPT2Backbone": get_gpt2_config,
-    # Add for future models, e.g., "MistralBackbone": get_mistral_config
+    "QwenBackbone": get_qwen_config,
 }
 
 MODEL_EXPORTERS = {
     "GemmaBackbone": get_gemma_weights_map,
     "GPT2Backbone": get_gpt2_weights_map,
-    # Add for future models, e.g., "MistralBackbone": get_mistral_weights_map
+    "QwenBackbone": get_qwen_weights_map,
 }
 
 MODEL_TOKENIZER_CONFIGS = {
     "GemmaTokenizer": get_gemma_tokenizer_config,
     "GPT2Tokenizer": get_gpt2_tokenizer_config,
-    # Add for future models, e.g., "MistralTokenizer":
-    # get_mistral_tokenizer_config
+    "QwenTokenizer": get_qwen_tokenizer_config,
 }
 
 
@@ -62,50 +71,64 @@ def export_backbone(backbone, path, include_lm_head=False):
     weights_dict = get_weights_fn(backbone, include_lm_head=include_lm_head)
     if not weights_dict:
         raise ValueError("No weights to save.")
+
     # Save config
     os.makedirs(path, exist_ok=True)
     config_path = os.path.join(path, "config.json")
+
+    # Handle Config Objects (GPT2/Qwen) vs Dicts (Gemma)
+    config_to_save = hf_config
+    if hasattr(hf_config, "to_dict"):
+        config_to_save = hf_config.to_dict()
+
     with open(config_path, "w") as f:
-        json.dump(hf_config.to_dict(), f)
+        json.dump(config_to_save, f, indent=2)
+
     # Save weights based on backend
     weights_path = os.path.join(path, "model.safetensors")
     if backend == "torch":
+        # Lazy import to prevent crash on TF-only environments
         import torch
         from safetensors.torch import save_file
 
         weights_dict_torch = {}
-
         for k, v in weights_dict.items():
             tensor = v.value if hasattr(v, "value") else v
 
-            # Torch tensor -> move to CPU
             if isinstance(tensor, torch.Tensor):
                 t = tensor.detach().to("cpu")
-
-            # TensorFlow / JAX -> convert via numpy()
             elif hasattr(tensor, "numpy"):
                 t = torch.tensor(tensor.numpy())
-
-            # numpy array
             elif hasattr(tensor, "__array__"):
                 t = torch.tensor(tensor)
-
             else:
-                raise TypeError(f"Unsupported tensor type: {type(tensor)}")
+                t = tensor
+
+            if hasattr(t, "contiguous"):
+                t = t.contiguous()
 
-            weights_dict_torch[k] = t.contiguous()
+            weights_dict_torch[k] = t
 
-        # ----  GPT-2 tied weights ----
+        # Handle Tied Weights (GPT-2, Qwen)
+        # Safetensors crashes if we try to save the same shared memory twice.
         if (
+            "lm_head.weight" in weights_dict_torch
+            and "model.embed_tokens.weight" in weights_dict_torch
+        ):
+            # Qwen / Llama naming convention
+            wte = weights_dict_torch["model.embed_tokens.weight"]
+            lm = weights_dict_torch["lm_head.weight"]
+            if wte.data_ptr() == lm.data_ptr():
+                weights_dict_torch["lm_head.weight"] = lm.clone().contiguous()
+        elif (
             "lm_head.weight" in weights_dict_torch
             and "transformer.wte.weight" in weights_dict_torch
         ):
+            # GPT-2 naming convention
             wte = weights_dict_torch["transformer.wte.weight"]
             lm = weights_dict_torch["lm_head.weight"]
-
-        if wte.data_ptr() == lm.data_ptr():
-            weights_dict_torch["lm_head.weight"] = lm.clone().contiguous()
-        # --------------------------------
+            if wte.data_ptr() == lm.data_ptr():
+                weights_dict_torch["lm_head.weight"] = lm.clone().contiguous()
 
         save_file(weights_dict_torch, weights_path, metadata={"format": "pt"})
 
@@ -129,46 +152,41 @@ def export_tokenizer(tokenizer, path):
         path: str. Path to save the exported tokenizer.
     """
     os.makedirs(path, exist_ok=True)
+
     # Save tokenizer assets
+    # BytePairTokenizer (GPT2, Qwen) -> "vocabulary.json", "merges.txt"
+    # SentencePieceTokenizer (Gemma) -> "vocabulary.spm"
     tokenizer.save_assets(path)
+
     # Export tokenizer config
     tokenizer_type = tokenizer.__class__.__name__
     if tokenizer_type not in MODEL_TOKENIZER_CONFIGS:
         raise ValueError(
-            "Export to Transformers format not implemented for {tokenizer_type}"
+            f"Export to Transformer format not implemented for {tokenizer_type}"
         )
     get_tokenizer_config_fn = MODEL_TOKENIZER_CONFIGS[tokenizer_type]
     tokenizer_config = get_tokenizer_config_fn(tokenizer)
     tokenizer_config_path = os.path.join(path, "tokenizer_config.json")
     with open(tokenizer_config_path, "w") as f:
         json.dump(tokenizer_config, f, indent=4)
 
+    # Rename files to match Hugging Face expectations
     if tokenizer_type == "GemmaTokenizer":
-        # Rename vocabulary file
         vocab_spm_path = os.path.join(path, "vocabulary.spm")
         tokenizer_model_path = os.path.join(path, "tokenizer.model")
         if os.path.exists(vocab_spm_path):
             shutil.move(vocab_spm_path, tokenizer_model_path)
         else:
-            warnings.warn(
-                f"{vocab_spm_path} not found. Tokenizer may not load "
-                "correctly. Ensure that the tokenizer configuration "
-                "is correct and that the vocabulary file is present "
-                "in the original model."
-            )
-    elif tokenizer_type == "GPT2Tokenizer":
-        # Rename vocabulary file
+            warnings.warn(f"{vocab_spm_path} not found.")
+
+    elif tokenizer_type in ["GPT2Tokenizer", "QwenTokenizer"]:
+        # Both GPT-2 and Qwen (BPE) use vocab.json in HF
         vocab_json_path = os.path.join(path, "vocabulary.json")
-        renamed_vocab_json_path = os.path.join(path, "vocab.json")
+        vocab_hf_path = os.path.join(path, "vocab.json")
         if os.path.exists(vocab_json_path):
-            shutil.move(vocab_json_path, renamed_vocab_json_path)
+            shutil.move(vocab_json_path, vocab_hf_path)
         else:
-            warnings.warn(
-                f"{vocab_json_path} not found. Tokenizer may not load "
-                "correctly. Ensure that the tokenizer configuration "
-                "is correct and that the vocabulary file is present "
-                "in the original model."
-            )
+            warnings.warn(f"{vocab_json_path} not found.")
 
 
 def export_to_safetensors(keras_model, path):
diff --git a/keras_hub/src/utils/transformers/export/qwen.py b/keras_hub/src/utils/transformers/export/qwen.py
@@ -0,0 +1,150 @@
+import keras.ops as ops
+import transformers
+
+
+def get_qwen_config(backbone):
+    """Convert Keras Qwen config to Hugging Face Qwen2Config."""
+    # Qwen2Config handles the architecture specifics (RoPE, RMSNorm, SwiGLU)
+    return transformers.Qwen2Config(
+        vocab_size=backbone.vocabulary_size,
+        hidden_size=backbone.hidden_dim,
+        num_hidden_layers=backbone.num_layers,
+        num_attention_heads=backbone.num_query_heads,
+        num_key_value_heads=backbone.num_key_value_heads,
+        intermediate_size=backbone.intermediate_dim,
+        hidden_act="silu",  # Qwen uses SwiGLU (SiLU activation)
+        rms_norm_eps=backbone.layer_norm_epsilon,
+        rope_theta=backbone.rope_max_wavelength,
+        tie_word_embeddings=backbone.tie_word_embeddings,
+        # Default initialization parameters
+        initializer_range=0.02,
+        use_cache=True,
+    )
+
+
+def get_qwen_weights_map(backbone, include_lm_head=False):
+    """Create a weights map for a given Qwen model."""
+    weights_map = {}
+
+    # 1. Embeddings
+    # Keras: token_embedding.embeddings
+    # HF: model.embed_tokens.weight
+    weights_map["model.embed_tokens.weight"] = backbone.get_layer(
+        "token_embedding"
+    ).embeddings
+
+    for i in range(backbone.num_layers):
+        # Access the decoder layer
+        decoder_layer = backbone.get_layer(f"transformer_layer_{i}")
+
+        # --- Normalization ---
+        # Input Norm (Pre-Attention)
+        # Keras uses 'scale' (gamma), HF uses 'weight'
+        weights_map[f"model.layers.{i}.input_layernorm.weight"] = (
+            decoder_layer._self_attention_layernorm.scale
+        )
+
+        # Post Attention Norm (Pre-MLP)
+        weights_map[f"model.layers.{i}.post_attention_layernorm.weight"] = (
+            decoder_layer._feedforward_layernorm.scale
+        )
+
+        # --- Attention ---
+        # QwenAttention uses EinsumDense for Q/K/V/O
+        # Keras Shape: (hidden_dim, num_heads, head_dim)
+        # HF Shape:    (num_heads * head_dim, hidden_dim) -> Transposed Linear
+
+        attn_layer = decoder_layer._self_attention_layer
+
+        # Query
+        q_kernel = attn_layer._query_dense.kernel
+        q_kernel = ops.reshape(q_kernel, (backbone.hidden_dim, -1))
+        weights_map[f"model.layers.{i}.self_attn.q_proj.weight"] = (
+            ops.transpose(q_kernel)
+        )
+
+        # Keras: (num_heads, head_dim) -> HF: (hidden_dim,)
+        weights_map[f"model.layers.{i}.self_attn.q_proj.bias"] = ops.reshape(
+            attn_layer._query_dense.bias, (-1,)
+        )
+
+        # Key
+        k_kernel = attn_layer._key_dense.kernel
+        k_kernel = ops.reshape(k_kernel, (backbone.hidden_dim, -1))
+        weights_map[f"model.layers.{i}.self_attn.k_proj.weight"] = (
+            ops.transpose(k_kernel)
+        )
+
+        weights_map[f"model.layers.{i}.self_attn.k_proj.bias"] = ops.reshape(
+            attn_layer._key_dense.bias, (-1,)
+        )
+
+        # Value
+        v_kernel = attn_layer._value_dense.kernel
+        v_kernel = ops.reshape(v_kernel, (backbone.hidden_dim, -1))
+        weights_map[f"model.layers.{i}.self_attn.v_proj.weight"] = (
+            ops.transpose(v_kernel)
+        )
+
+        weights_map[f"model.layers.{i}.self_attn.v_proj.bias"] = ops.reshape(
+            attn_layer._value_dense.bias, (-1,)
+        )
+
+        # Output
+        o_kernel = attn_layer._output_dense.kernel
+        o_kernel = ops.reshape(o_kernel, (-1, backbone.hidden_dim))
+        weights_map[f"model.layers.{i}.self_attn.o_proj.weight"] = (
+            ops.transpose(o_kernel)
+        )
+
+        # --- MLP (SwiGLU) ---
+        # Gate (feedforward_gate_dense)
+        gate_kernel = decoder_layer._feedforward_gate_dense.kernel
+        weights_map[f"model.layers.{i}.mlp.gate_proj.weight"] = ops.transpose(
+            gate_kernel
+        )
+
+        # Up (feedforward_intermediate_dense)
+        up_kernel = decoder_layer._feedforward_intermediate_dense.kernel
+        weights_map[f"model.layers.{i}.mlp.up_proj.weight"] = ops.transpose(
+            up_kernel
+        )
+
+        # Down (feedforward_output_dense)
+        down_kernel = decoder_layer._feedforward_output_dense.kernel
+        weights_map[f"model.layers.{i}.mlp.down_proj.weight"] = ops.transpose(
+            down_kernel
+        )
+
+    # Final Norm
+    weights_map["model.norm.weight"] = backbone.get_layer(
+        "sequence_output_layernorm"
+    ).scale
+
+    # LM Head
+    if include_lm_head:
+        if backbone.tie_word_embeddings:
+            # If tied, point to embeddings (Exporter handles cloning)
+            weights_map["lm_head.weight"] = weights_map[
+                "model.embed_tokens.weight"
+            ]
+        else:
+            # If not tied, QwenBackbone uses ReversibleEmbedding.
+            lm_head_w = backbone.get_layer("token_embedding").reverse_embeddings
+            weights_map["lm_head.weight"] = ops.transpose(lm_head_w)
+
+    return weights_map
+
+
+def get_qwen_tokenizer_config(tokenizer):
+    """Convert Keras Qwen tokenizer config to Hugging Face."""
+    # Qwen2 uses BPE. We specify the class and basic special tokens.
+    # The actual vocab/merges files are handled by the exporter.
+    return {
+        "tokenizer_class": "Qwen2Tokenizer",
+        "bos_token": None,  # Qwen often doesn't use BOS
+        "eos_token": "<|endoftext|>",
+        "pad_token": "<|endoftext|>",  # Often mapped to EOS or null
+        "unk_token": None,
+        "model_max_length": 32768,  # Default window size
+    }
diff --git a/keras_hub/src/utils/transformers/export/qwen_test.py b/keras_hub/src/utils/transformers/export/qwen_test.py
