fix conversion error when processing tensorblock-TinyStories-656K GGUF

fanfu/gguf_to_hf.py

@@ -7,6 +7,7 @@
 
 import json
 import logging
+import re
 import struct
 import sys
 from pathlib import Path
@@ -112,6 +113,42 @@ def dequantize_tensor(tensor) -> np.ndarray:
     return data
 
 
+GGUF_TO_HF_TENSOR_PATTERNS = [
+    (r"^token_embd\.weight$", "model.embed_tokens.weight"),
+    (r"^output_norm\.weight$", "model.norm.weight"),
+    (r"^blk\.(\d+)\.attn_norm\.weight$", r"model.layers.\1.input_layernorm.weight"),
+    (r"^blk\.(\d+)\.attn_q\.weight$", r"model.layers.\1.self_attn.q_proj.weight"),
+    (r"^blk\.(\d+)\.attn_k\.weight$", r"model.layers.\1.self_attn.k_proj.weight"),
+    (r"^blk\.(\d+)\.attn_v\.weight$", r"model.layers.\1.self_attn.v_proj.weight"),
+    (r"^blk\.(\d+)\.attn_output\.weight$", r"model.layers.\1.self_attn.o_proj.weight"),
+    (r"^blk\.(\d+)\.ffn_norm\.weight$", r"model.layers.\1.post_attention_layernorm.weight"),
+    (r"^blk\.(\d+)\.ffn_gate\.weight$", r"model.layers.\1.mlp.gate_proj.weight"),
+    (r"^blk\.(\d+)\.ffn_up\.weight$", r"model.layers.\1.mlp.up_proj.weight"),
+    (r"^blk\.(\d+)\.ffn_down\.weight$", r"model.layers.\1.mlp.down_proj.weight"),
+]
+
+GGUF_TO_HF_CONFIG_MAP = {
+    "block_count": "num_hidden_layers",
+    "context_length": "max_position_embeddings",
+    "embedding_length": "hidden_size",
+    "feed_forward_length": "intermediate_size",
+    "attention.head_count": "num_attention_heads",
+    "attention.head_count_kv": "num_key_value_heads",
+    "attention.layer_norm_rms_epsilon": "rms_norm_eps",
+    "rope.freq_base": "rope_theta",
+    "rope.dimension_count": None,  # not needed in HF config
+}
+
+
+def map_tensor_name_gguf_to_hf(name: str, state_dict: dict[str, Any]) -> str:
+    """Map a GGUF tensor name to HuggingFace naming convention."""
+    for pattern, replacement in GGUF_TO_HF_TENSOR_PATTERNS:
+        m = re.match(pattern, name)
+        if m:
+            return m.expand(replacement)
+    return name
+
+
 def extract_tokenizer_from_gguf(gguf_path: Path, output_dir: Path) -> None:
     """Extract full tokenizer from GGUF and create HF-compatible tokenizer files."""
     import gguf
@@ -255,14 +292,10 @@ def convert_gguf_to_hf(
         if key.startswith("general."):
             short_key = key[len("general."):]
             if short_key == "architecture":
-                arch = field.parts[0]
-                if hasattr(reader, 'tensors'):
-                    try:
-                        config["architectures"] = [f"{arch}ForCausalLM"]
-                        config["model_type"] = arch
-                    except Exception:
-                        config["architectures"] = [str(field.contents())]
-                        config["model_type"] = str(field.contents())
+                arch = str(field.contents())
+                arch_capitalized = arch[0].upper() + arch[1:] if arch else arch
+                config["architectures"] = [f"{arch_capitalized}ForCausalLM"]
+                config["model_type"] = arch
             elif short_key == "name":
                 config["name"] = str(field.contents())
 
@@ -273,10 +306,18 @@ def convert_gguf_to_hf(
                 val = field.contents()
                 if isinstance(val, np.ndarray):
                     val = val.tolist()
-                config[short_key] = val
+                hf_key = GGUF_TO_HF_CONFIG_MAP.get(short_key, short_key)
+                if hf_key is not None:
+                    config[hf_key] = val
             except Exception:
                 pass
 
+    # Detect tie_word_embeddings
+    tensor_names = [t.name for t in reader.tensors]
+    has_output_weight = any(n == "output.weight" for n in tensor_names)
+    if not has_output_weight and "vocab_size" in config:
+        config["tie_word_embeddings"] = True
+
     with open(output / "config.json", "w") as f:
         json.dump(config, f, indent=2, ensure_ascii=False)
     logger.info("Saved config.json")
@@ -307,14 +348,21 @@ def convert_gguf_to_hf(
 
     logger.info(f"Loaded {len(state_dict)} tensors, skipped {len(skipped)}")
 
+    logger.info("Mapping tensor names to HuggingFace conventions...")
+    mapped_state: dict[str, np.ndarray] = {}
+    for n, d in state_dict.items():
+        new_name = map_tensor_name_gguf_to_hf(n, mapped_state)
+        mapped_state[new_name] = d
+    logger.info(f"Remapped {len(mapped_state)} tensors")
+
     try:
         from safetensors.torch import save_file
         import torch
     except ImportError:
         return ToolResult(success=False, error="safetensors or torch not installed. Run: pip install safetensors torch")
 
     torch_state = {}
-    for n, d in state_dict.items():
+    for n, d in mapped_state.items():
         torch_state[n] = torch.from_numpy(d.copy())
 
     total = sum(t.numel() * t.element_size() for t in torch_state.values())
@@ -329,5 +377,5 @@ def convert_gguf_to_hf(
 
     return ToolResult(
         success=True,
-        data={"output_dir": str(output), "tensors": len(state_dict), "skipped": len(skipped)},
+        data={"output_dir": str(output), "tensors": len(mapped_state), "skipped": len(skipped)},
     )