feat(tools): HuggingFace config parser for transformer spec

tools/hf_parser.py

@@ -0,0 +1,365 @@
+#!/usr/bin/env python3
+"""Parse HuggingFace model config.json into ModelPack transformer spec format.
+
+This tool maps HuggingFace Transformers config.json fields to the ModelPack
+unified transformer specification vocabulary defined in PR #111
+(docs/architecture.md by @aftersnow).
+
+Usage:
+    python tools/hf_parser.py meta-llama/Meta-Llama-3-8B
+    python tools/hf_parser.py mistralai/Mistral-7B-v0.3
+    python tools/hf_parser.py --file path/to/config.json
+
+The output is a YAML spec file following the ModelPack transformer spec format.
+Fields that cannot be reliably inferred from config.json are marked as
+NEEDS_REVIEW for human verification.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from pathlib import Path
+
+NEEDS_REVIEW = "__NEEDS_REVIEW__"
+
+# Maps HuggingFace config.json field names to ModelPack transformer spec paths.
+# Based on PR #111's field vocabulary (docs/architecture.md).
+FIELD_MAP = {
+    # Top-level transformer fields
+    "vocab_size": "vocabulary_size",
+    "hidden_size": "hidden_size",
+    # Position embedding
+    "max_position_embeddings": "position_embedding.max_position_embeddings",
+    "rope_theta": "position_embedding.rope_theta",
+    "rope_scaling": "position_embedding.rope_scaling",
+    # Attention
+    "num_attention_heads": "attention.num_attention_heads",
+    "num_key_value_heads": "attention.num_key_value_heads",
+    "head_dim": "attention.head_dim",
+    # FFN / MLP
+    "intermediate_size": "mlp.intermediate_size",
+    # Transformer layers
+    "num_hidden_layers": "num_layers",
+    # Normalization
+    "rms_norm_eps": "norm.epsilon",
+    # MoE fields
+    "num_local_experts": "moe.num_experts",
+    "num_experts_per_tok": "moe.top_k",
+    "num_experts": "moe.num_experts",
+    "n_routed_experts": "moe.num_experts",  # DeepSeek naming variant
+    # MLA fields (DeepSeek)
+    "kv_lora_rank": "attention.kv_lora_rank",
+    "q_lora_rank": "attention.q_lora_rank",
+    "qk_nope_head_dim": "attention.qk_nope_head_dim",
+    "qk_rope_head_dim": "attention.qk_rope_head_dim",
+    "v_head_dim": "attention.v_head_dim",
+}
+
+# Known model type → attention type mapping
+ATTENTION_TYPE_MAP = {
+    "llama": "gqa",
+    "mistral": "gqa",
+    "mixtral": "gqa",
+    "qwen2": "gqa",
+    "qwen2_moe": "gqa",
+    "gemma": "gqa",
+    "gemma2": "gqa",
+    "phi3": "gqa",
+    "deepseek_v2": "mla",
+    "deepseek_v3": "mla",
+    "gpt2": "mha",
+    "gpt_neo": "mha",
+    "gpt_neox": "mha",
+    "falcon": "mha",
+}
+
+# Known model type → FFN type mapping
+FFN_TYPE_MAP = {
+    "llama": "mlp",
+    "mistral": "mlp",
+    "mixtral": "moe",
+    "qwen2": "mlp",
+    "qwen2_moe": "moe",
+    "gemma": "mlp",
+    "gemma2": "mlp",
+    "phi3": "mlp",
+    "deepseek_v2": "moe",
+    "deepseek_v3": "moe",
+    "gpt2": "mlp",
+    "gpt_neo": "mlp",
+    "gpt_neox": "mlp",
+    "falcon": "mlp",
+}
+
+# Known model type → activation function mapping
+ACTIVATION_MAP = {
+    "llama": "silu",
+    "mistral": "silu",
+    "mixtral": "silu",
+    "qwen2": "silu",
+    "qwen2_moe": "silu",
+    "gemma": "gelu",
+    "gemma2": "gelu",
+    "phi3": "silu",
+    "gpt2": "gelu",
+    "gpt_neo": "gelu",
+    "gpt_neox": "gelu",
+    "falcon": "gelu",
+}
+
+
+def _set_nested(d: dict, path: str, value) -> None:
+    """Set a value in a nested dict using dot-separated path."""
+    keys = path.split(".")
+    for key in keys[:-1]:
+        d = d.setdefault(key, {})
+    d[keys[-1]] = value
+
+
+def _get_nested(d: dict, path: str, default=None):
+    """Get a value from a nested dict using dot-separated path."""
+    keys = path.split(".")
+    for key in keys:
+        if not isinstance(d, dict) or key not in d:
+            return default
+        d = d[key]
+    return d
+
+
+def parse_hf_config(raw: dict) -> dict:
+    """Parse a HuggingFace config.json dict into ModelPack transformer spec.
+
+    Args:
+        raw: The parsed config.json dict from HuggingFace.
+
+    Returns:
+        A dict following the ModelPack transformer spec format.
+    """
+    result: dict = {}
+    model_type = raw.get("model_type", "").lower()
+
+    # Map static fields
+    for hf_key, mp_path in FIELD_MAP.items():
+        if hf_key in raw and raw[hf_key] is not None:
+            _set_nested(result, mp_path, raw[hf_key])
+
+    # Derive head_dim if absent
+    if "attention" in result and "head_dim" not in result.get("attention", {}):
+        hidden = result.get("hidden_size")
+        n_heads = _get_nested(result, "attention.num_attention_heads")
+        if hidden and n_heads:
+            _set_nested(result, "attention.head_dim", hidden // n_heads)
+
+    # Set architecture type
+    result["type"] = "decoder"
+    result["architecture_version"] = "0.1.0"
+
+    # Infer attention type from model_type
+    attn_type = ATTENTION_TYPE_MAP.get(model_type, NEEDS_REVIEW)
+    _set_nested(result, "attention.type", attn_type)
+    _set_nested(result, "attention.is_causal", True)
+
+    # Check for sliding window attention
+    if raw.get("sliding_window") is not None:
+        _set_nested(result, "attention.sliding_window", raw["sliding_window"])
+
+    # Infer FFN type
+    ffn_type = FFN_TYPE_MAP.get(model_type, NEEDS_REVIEW)
+    result["ffn_type"] = ffn_type
+
+    # Set activation function
+    hf_activation = raw.get("hidden_act", raw.get("activation_function"))
+    if hf_activation:
+        activation = hf_activation.lower()
+        if "silu" in activation or "swish" in activation:
+            activation = "silu"
+        elif "gelu" in activation:
+            activation = "gelu"
+        elif "relu" in activation:
+            activation = "relu"
+    else:
+        activation = ACTIVATION_MAP.get(model_type, NEEDS_REVIEW)
+
+    if ffn_type == "mlp":
+        _set_nested(result, "mlp.activation", activation)
+        # Most modern models use gated activation (SwiGLU, GeGLU)
+        use_gated = model_type in (
+            "llama", "mistral", "mixtral", "qwen2", "qwen2_moe", "phi3",
+            "gemma", "gemma2", "deepseek_v2", "deepseek_v3",
+        )
+        _set_nested(result, "mlp.use_gated_activation", use_gated)
+    elif ffn_type == "moe":
+        _set_nested(result, "moe.activation", activation)
+        # MoE-specific fields
+        if "moe_intermediate_size" in raw:
+            _set_nested(result, "moe.moe_intermediate_size", raw["moe_intermediate_size"])
+        if "num_shared_experts" in raw:
+            _set_nested(result, "moe.num_shared_experts", raw["num_shared_experts"])
+        if "shared_expert_intermediate_size" in raw:
+            _set_nested(
+                result, "moe.shared_expert_intermediate_size",
+                raw["shared_expert_intermediate_size"],
+            )
+        # DeepSeek MoE-specific fields (from PR #185 research)
+        if "routed_scaling_factor" in raw:
+            _set_nested(result, "moe.routed_scaling_factor", raw["routed_scaling_factor"])
+        if "topk_method" in raw:
+            _set_nested(result, "moe.topk_method", raw["topk_method"])
+        if "norm_topk_prob" in raw:
+            _set_nested(result, "moe.norm_topk_prob", raw["norm_topk_prob"])
+
+    # Mixed layers support (DeepSeek uses dense layers before switching to MoE)
+    if "first_k_dense_replace" in raw and "moe_layer_freq" in raw:
+        result["layer_structure"] = "mixed"
+        _set_nested(result, "mixed_layers.first_k_dense_replace", raw["first_k_dense_replace"])
+        _set_nested(result, "mixed_layers.moe_layer_freq", raw["moe_layer_freq"])
+
+    # Normalization
+    norm_type = "rmsnorm"  # Most modern models use RMSNorm
+    if model_type in ("gpt2", "gpt_neo"):
+        norm_type = "layernorm"
+    _set_nested(result, "norm.type", norm_type)
+
+    if "layer_norm_eps" in raw:
+        _set_nested(result, "norm.epsilon", raw["layer_norm_eps"])
+
+    # Tokenizer
+    _set_nested(result, "tokenizer.type", "bpe")
+    _set_nested(result, "tokenizer.library", "huggingface")
+
+    # Position embedding type
+    if model_type in ("gpt2", "gpt_neo"):
+        _set_nested(result, "position_embedding.type", "learned")
+    else:
+        _set_nested(result, "position_embedding.type", "rope")
+
+    # Embedding
+    tie_embeddings = raw.get("tie_word_embeddings", False)
+    _set_nested(result, "token_embedding.shared_embedding", tie_embeddings)
+
+    # Bias flags
+    attn_bias = raw.get("attention_bias", False)
+    _set_nested(result, "attention.has_qkv_bias", attn_bias)
+    _set_nested(result, "attention.has_output_bias", attn_bias)
+
+    mlp_bias = raw.get("mlp_bias", False)
+    if ffn_type == "mlp":
+        _set_nested(result, "mlp.has_bias", mlp_bias)
+
+    return result
+
+
+def format_yaml(spec: dict, indent: int = 0) -> str:
+    """Format a spec dict as YAML string."""
+    lines = []
+    prefix = "  " * indent
+    for key, value in spec.items():
+        if isinstance(value, dict):
+            lines.append(f"{prefix}{key}:")
+            lines.append(format_yaml(value, indent + 1))
+        elif isinstance(value, bool):
+            lines.append(f"{prefix}{key}: {str(value).lower()}")
+        elif isinstance(value, str):
+            if value == NEEDS_REVIEW:
+                lines.append(f"{prefix}{key}: {value}  # requires human review")
+            else:
+                lines.append(f'{prefix}{key}: "{value}"')
+        elif value is None:
+            lines.append(f"{prefix}{key}: null")
+        else:
+            lines.append(f"{prefix}{key}: {value}")
+    return "\n".join(lines)
+
+
+def load_config(source: str) -> dict:
+    """Load a config.json from a file path or HuggingFace model ID.
+
+    Args:
+        source: Either a local file path or a HuggingFace model ID.
+
+    Returns:
+        The parsed config.json dict.
+    """
+    path = Path(source)
+    if path.is_file():
+        with path.open(encoding="utf-8") as f:
+            return json.load(f)
+
+    # Try loading from HuggingFace Hub
+    try:
+        from huggingface_hub import hf_hub_download
+
+        config_path = hf_hub_download(repo_id=source, filename="config.json")
+        with open(config_path, encoding="utf-8") as f:
+            return json.load(f)
+    except ImportError:
+        print(
+            "error: huggingface_hub not installed. "
+            "Install with: pip install huggingface_hub",
+            file=sys.stderr,
+        )
+        sys.exit(1)
+    except Exception as e:
+        print(f"error: failed to load config from '{source}': {e}", file=sys.stderr)
+        sys.exit(1)
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(
+        description="Parse HuggingFace config.json into ModelPack transformer spec",
+    )
+    parser.add_argument(
+        "model",
+        help="HuggingFace model ID (e.g., meta-llama/Meta-Llama-3-8B) "
+        "or path to config.json",
+    )
+    parser.add_argument(
+        "--format",
+        choices=["yaml", "json"],
+        default="yaml",
+        help="Output format (default: yaml)",
+    )
+
+    args = parser.parse_args()
+
+    raw = load_config(args.model)
+    spec = parse_hf_config(raw)
+
+    model_type = raw.get("model_type", "unknown")
+    model_name = raw.get("_name_or_path", args.model)
+
+    if args.format == "json":
+        print(json.dumps(spec, indent=2))
+    else:
+        print(f"# ModelPack Transformer Spec")
+        print(f"# Generated from: {model_name}")
+        print(f"# Model type: {model_type}")
+        print(f"# NOTE: Fields marked NEEDS_REVIEW require human verification")
+        print()
+        print(format_yaml(spec))
+
+    # Report coverage
+    needs_review = []
+    _find_needs_review(spec, "", needs_review)
+    if needs_review:
+        print(f"\n# --- Fields requiring review ({len(needs_review)}) ---")
+        for field in needs_review:
+            print(f"#   - {field}")
+
+    return 0
+
+
+def _find_needs_review(d: dict, prefix: str, result: list) -> None:
+    """Recursively find all NEEDS_REVIEW fields."""
+    for key, value in d.items():
+        path = f"{prefix}.{key}" if prefix else key
+        if isinstance(value, dict):
+            _find_needs_review(value, path, result)
+        elif value == NEEDS_REVIEW:
+            result.append(path)
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())