model-unfolder 0.2.0__py3-none-any.whl

model_unfolder/__init__.py

@@ -0,0 +1,58 @@
+"""model_unfolder — turn any HuggingFace transformer into a clear architecture diagram.
+
+Quick start in a Jupyter notebook::
+
+    from model_unfolder import unfold
+    unfold("moonshotai/Kimi-K2-Instruct")
+
+Outside Jupyter::
+
+    diagram = unfold(cfg)
+    diagram.save("kimi_k2.html")
+"""
+from .diagram import Diagram
+from .parser import config_to_ir
+from .ir import ModelIR, LayerSpec, AttentionSpec, FFNSpec, CrossLayerEdge
+from .params import estimate_params
+
+__version__ = "0.2.0"
+
+__all__ = [
+    "unfold",
+    "show",
+    "Diagram",
+    "ModelIR",
+    "LayerSpec",
+    "AttentionSpec",
+    "FFNSpec",
+    "CrossLayerEdge",
+    "config_to_ir",
+    "estimate_params",
+]
+
+
+def unfold(cfg_or_id, token=None) -> Diagram:
+    """Unfold a transformer into a renderable architecture diagram.
+
+    Parameters
+    ----------
+    cfg_or_id
+        A HuggingFace ``PretrainedConfig`` instance, a model ID string
+        (e.g. ``"moonshotai/Kimi-K2-Instruct"``), or a plain ``dict`` of
+        ``config.json`` contents.
+    token
+        Optional Hugging Face token used only when ``cfg_or_id`` is a model ID.
+        If omitted, ``HF_TOKEN`` and legacy Hugging Face token env vars are used
+        when present.
+
+    Returns
+    -------
+    Diagram
+        Renders inline in Jupyter; otherwise call ``.save()`` or ``.to_html()``.
+    """
+    ir = config_to_ir(cfg_or_id, token=token)
+    return Diagram(ir)
+
+
+# friendly alias
+show = unfold

model_unfolder/adapters/__init__.py

@@ -0,0 +1,15 @@
+"""Adapter registry. Order matters: more specific adapters come first."""
+from . import custom, diffusor, transformer
+
+ADAPTERS = [
+    *custom.ADAPTERS,
+    *diffusor.ADAPTERS,
+    *transformer.ADAPTERS,
+]
+
+
+def find_adapter(cfg):
+    for a in ADAPTERS:
+        if a.matches(cfg):
+            return a
+    return None

model_unfolder/adapters/custom/__init__.py

@@ -0,0 +1,8 @@
+"""Project-local and experimental adapters.
+
+Drop custom adapters here when a model needs bespoke parsing before it is
+general enough to move into a domain package such as ``transformer``.
+"""
+
+ADAPTERS = []
+

model_unfolder/adapters/diffusor/__init__.py

@@ -0,0 +1,8 @@
+"""Diffusor-family adapters.
+
+This package is intentionally empty for now.  It gives diffusion-style models
+their own adapter namespace instead of stretching transformer assumptions.
+"""
+
+ADAPTERS = []
+

model_unfolder/adapters/transformer/__init__.py

@@ -0,0 +1,5 @@
+"""Transformer-family adapters and reusable block vocabularies."""
+
+from . import families
+
+ADAPTERS = families.ADAPTERS

model_unfolder/adapters/transformer/assembly.py

@@ -0,0 +1,57 @@
+"""Assembly helpers for transformer-family adapters."""
+from __future__ import annotations
+
+from collections.abc import Iterable, Mapping
+from typing import Any
+
+from ...ir import AttentionSpec, FFNSpec, LayerSpec
+from .blocks import decoder_layer_blocks, decoder_only_render_spec
+
+
+def decoder_layer(
+    index: int,
+    attention: AttentionSpec,
+    ffn: FFNSpec,
+    hidden_size: int,
+    *,
+    extra_blocks: Iterable[dict] | None = None,
+) -> LayerSpec:
+    """Build a decoder layer from parsed specs plus optional reusable parts."""
+    blocks = decoder_layer_blocks(attention, ffn, hidden_size)
+    if extra_blocks:
+        blocks.extend(extra_blocks)
+    return LayerSpec(
+        index=index,
+        attention=attention,
+        ffn=ffn,
+        blocks=blocks,
+    )
+
+
+def decoder_extras(
+    vocab_size: int,
+    hidden_size: int,
+    tie_word_embeddings: bool,
+    *extra_maps: Mapping[str, Any] | None,
+) -> dict:
+    """Build top-level extras shared by decoder-only transformer models."""
+    extras = {
+        "render": decoder_only_render_spec(
+            vocab_size,
+            hidden_size,
+            tie_word_embeddings,
+        )
+    }
+    for extra in extra_maps:
+        if not extra:
+            continue
+        _merge_extras(extras, extra)
+    return extras
+
+
+def _merge_extras(target: dict, extra: Mapping[str, Any]) -> None:
+    for key, value in extra.items():
+        if key == "external_pathways" and key in target:
+            target[key].extend(value)
+        else:
+            target[key] = value

model_unfolder/adapters/transformer/blocks.py

@@ -0,0 +1,238 @@
+"""Reusable transformer block descriptions for renderers.
+
+Adapters attach these block parts to the IR.  Renderers can then draw generic
+decoder-only transformer layouts without re-discovering model-specific names
+or labels with another layer of ``if model_type`` logic.
+
+Each block carries two orthogonal tags:
+
+* ``role`` — semantic ("norm", "attention", "ffn", "residual", "gate") used
+  for tooltips, click handlers, and the inspect cards.
+* ``kind`` — rendering shape ("norm", "linear", "activation", "attention",
+  "ffn", "residual_add", "gate_mul", "embedding", "output", "source") used
+  by the architecture view to pick a glyph and lay out a slot.
+
+Edges between blocks travel on the destination side as plain string fields:
+
+* ``residual_from: "<other_block_id>"`` — the residual_add block consumes the
+  *input* of the named block (the standard pre-attention bypass pattern).
+* ``lane: "left" | "right"`` — the block is rendered off the central chain
+  and connected via ``tap_from`` / ``feeds``. Reusable parts such as
+  per-layer embeddings use this instead of model-specific renderer logic.
+"""
+from __future__ import annotations
+
+from ...labels import activation_label
+from ...ir import AttentionSpec, FFNSpec
+from .common import format_dim as _fmt
+
+
+def decoder_only_render_spec(vocab_size: int, hidden_size: int, tie_word_embeddings: bool) -> dict:
+    return {
+        "family": "transformer",
+        "layout": "decoder_only",
+        "model_blocks": decoder_model_blocks(vocab_size, hidden_size, tie_word_embeddings),
+    }
+
+
+def decoder_model_blocks(vocab_size: int, hidden_size: int, tie_word_embeddings: bool) -> list[dict]:
+    vocab = _fmt(vocab_size)
+    hidden = _fmt(hidden_size)
+    tied = " (tied with output)" if tie_word_embeddings else ""
+    return [
+        {
+            "id": "tok_text",
+            "role": "input",
+            "kind": "source",
+            "label": "Tokenized text",
+            "title": "Tokenized text",
+            "description": "Input token IDs; shape [batch, seq_len]",
+        },
+        {
+            "id": "embed",
+            "role": "embedding",
+            "kind": "embedding",
+            "label": "Token Embedding layer",
+            "title": "Token embedding",
+            "description": f"{vocab} x {hidden}{tied}",
+        },
+        {
+            "id": "final_rms",
+            "role": "norm",
+            "kind": "norm",
+            "label": "Final RMSNorm",
+            "title": "Final norm",
+            "description": f"RMSNorm; dim {hidden}",
+        },
+        {
+            "id": "lm_head",
+            "role": "output",
+            "kind": "output",
+            "label": "Linear output layer",
+            "title": "LM head",
+            "description": f"{hidden} -> {vocab}" + (" (tied)" if tie_word_embeddings else ""),
+        },
+    ]
+
+
+def decoder_layer_blocks(attention: AttentionSpec, ffn: FFNSpec, hidden_size: int) -> list[dict]:
+    hidden = _fmt(hidden_size)
+    return [
+        {
+            "id": "rms1",
+            "role": "norm",
+            "kind": "norm",
+            "label": "RMSNorm",
+            "title": "Pre-attention norm",
+            "description": f"RMSNorm; dim {hidden}",
+        },
+        {
+            "id": "attn",
+            "role": "attention",
+            "kind": "attention",
+            "label": attention_label(attention),
+            "title": attention_title(attention),
+            "description": describe_attention(attention),
+        },
+        {
+            "id": "add1",
+            "role": "residual",
+            "kind": "residual_add",
+            "residual_from": "rms1",
+            "label": "+",
+            "title": "Residual add",
+            "description": "block input + attention output",
+        },
+        {
+            "id": "rms2",
+            "role": "norm",
+            "kind": "norm",
+            "label": "RMSNorm",
+            "title": "Pre-FFN norm",
+            "description": f"RMSNorm; dim {hidden}",
+        },
+        {
+            "id": "ffn",
+            "role": "ffn",
+            "kind": "ffn",
+            "label": "MoE" if ffn.kind == "moe" else "Feed-Forward",
+            "title": "Mixture of experts" if ffn.kind == "moe" else "Feed-forward",
+            "description": describe_ffn(ffn),
+            "detail_view": "moe" if ffn.kind == "moe" else "gated_ffn",
+            "children": ffn_child_blocks(ffn, hidden_size),
+        },
+        {
+            "id": "add2",
+            "role": "residual",
+            "kind": "residual_add",
+            "residual_from": "rms2",
+            "label": "+",
+            "title": "Residual add",
+            "description": "post-attention + FFN output",
+        },
+    ]
+
+
+def ffn_child_blocks(ffn: FFNSpec, hidden_size: int) -> list[dict]:
+    hidden = _fmt(hidden_size)
+    inter = _fmt(ffn.expert_intermediate_size or ffn.intermediate_size)
+    activation = activation_label(ffn.activation)
+    children = [
+        {
+            "id": "gate_proj",
+            "title": "Gate projection",
+            "description": f"Linear; {hidden} -> {inter} (gated path through {activation})",
+        },
+        {"id": "up_proj", "title": "Up projection", "description": f"Linear; {hidden} -> {inter}"},
+        {
+            "id": "silu",
+            "title": f"{activation} activation",
+            "description": "Element-wise non-linearity applied to the gate path",
+        },
+        {
+            "id": "mul",
+            "title": "Element-wise multiply",
+            "description": f"{activation}(gate) x up; combines the gated and ungated paths",
+        },
+        {"id": "down_proj", "title": "Down projection", "description": f"Linear; {inter} -> {hidden}"},
+    ]
+    if ffn.kind == "moe":
+        n_experts = _fmt(ffn.num_experts) if ffn.num_experts else "N"
+        n_active = ffn.num_experts_per_tok or "k"
+        n_shared = ffn.num_shared_experts or 0
+        expert_desc = (
+            f"Dense FFN; {hidden} -> {inter} -> {hidden}; "
+            f"only top-{n_active} of {n_experts} active per token"
+            + (f"; plus {n_shared} shared expert(s) always active" if n_shared else "")
+        )
+        children.extend(
+            [
+                {
+                    "id": "router",
+                    "title": "Router",
+                    "description": f"Linear; {hidden} -> {n_experts} (selects top-{n_active} experts per token)",
+                },
+                {"id": "expert_1", "title": "Expert FFN", "description": expert_desc},
+                {"id": "expert_k", "title": "Expert FFN", "description": expert_desc},
+                {"id": "expert_kp1", "title": "Expert FFN", "description": expert_desc},
+                {"id": "expert_n", "title": "Expert FFN", "description": expert_desc},
+                {
+                    "id": "add_moe",
+                    "title": "Weighted sum",
+                    "description": f"Combines top-{n_active} expert outputs weighted by router probabilities",
+                },
+            ]
+        )
+    return children
+
+
+def attention_label(attention: AttentionSpec) -> list[str]:
+    kind = attention.kind
+    if kind == "mla":
+        return ["Multi-Head Latent", "Attention"]
+    if kind == "gqa":
+        return ["Grouped-Query", "Attention"]
+    if kind == "mqa":
+        return ["Multi-Query", "Attention"]
+    return ["Multi-Head", "Attention"]
+
+
+def attention_title(attention: AttentionSpec) -> str:
+    titles = {
+        "mla": "Multi-head latent attention",
+        "gqa": "Grouped-query attention",
+        "mqa": "Multi-query attention",
+    }
+    return titles.get(attention.kind, "Attention")
+
+
+def describe_attention(attention: AttentionSpec) -> str:
+    if attention.kind == "mla":
+        text = (
+            f"Multi-head latent attention; {attention.num_heads} heads; "
+            f"KV LoRA {_fmt(attention.kv_lora_rank)}"
+        )
+        if attention.q_lora_rank:
+            text += f"; Q LoRA {_fmt(attention.q_lora_rank)}"
+        return text
+    if attention.kind == "gqa":
+        return (
+            f"Grouped-query; {attention.num_heads} Q / {attention.num_kv_heads} KV heads; "
+            f"head dim {_fmt(attention.head_dim)}"
+        )
+    if attention.kind == "mqa":
+        return f"Multi-query; {attention.num_heads} Q / 1 KV head"
+    return f"Multi-head; {attention.num_heads} heads; head dim {_fmt(attention.head_dim)}"
+
+
+def describe_ffn(ffn: FFNSpec) -> str:
+    if ffn.kind == "moe":
+        text = f"MoE; {_fmt(ffn.num_experts)} experts; top-{ffn.num_experts_per_tok}"
+        if ffn.num_shared_experts:
+            text += f" + {ffn.num_shared_experts} shared"
+        if ffn.num_experts and ffn.num_experts_per_tok:
+            text += f"; {100 * ffn.num_experts_per_tok / ffn.num_experts:.1f}% active"
+        text += f"; expert hidden {_fmt(ffn.expert_intermediate_size or ffn.intermediate_size)}"
+        return text
+    gated = "gated " if ffn.gated else ""
+    return f"{gated}FFN; {activation_label(ffn.activation)}; hidden {_fmt(ffn.intermediate_size)}"

model_unfolder/adapters/transformer/common.py

@@ -0,0 +1,35 @@
+"""Shared helpers for transformer-family config adapters."""
+from __future__ import annotations
+
+from typing import Any
+
+
+def get_config_value(cfg: Any, name: str, default=None):
+    """Get a config value from a dict or a HuggingFace config object."""
+    if isinstance(cfg, dict):
+        return cfg.get(name, default)
+    return getattr(cfg, name, default)
+
+
+def architecture_name(cfg: Any, fallback: str) -> str:
+    architectures = get_config_value(cfg, "architectures") or []
+    return architectures[0] if architectures else get_config_value(cfg, "model_type", fallback)
+
+
+def model_name(cfg: Any, fallback: str) -> str:
+    name = (
+        get_config_value(cfg, "_name_or_path")
+        or get_config_value(cfg, "name_or_path")
+        or fallback
+    )
+    return str(name).split("/")[-1] if name else fallback
+
+
+def format_dim(value: Any) -> str:
+    """Human-readable dimension text for adapter-authored metadata."""
+    if value is None:
+        return "?"
+    try:
+        return f"{int(value):,}"
+    except (TypeError, ValueError):
+        return str(value)

model_unfolder/adapters/transformer/families/__init__.py

@@ -0,0 +1,12 @@
+"""Transformer model-family adapters.
+
+These modules translate family-specific HuggingFace config dialects into the
+shared transformer IR pieces in ``model_unfolder.adapters.transformer``.
+"""
+
+from . import deepseek, gemma4, llama
+
+# Order matters: more specific adapters first.  ``gemma4`` claims its own
+# top-level ``model_type`` / architecture and must run before the generic
+# llama-family matcher.
+ADAPTERS = [deepseek, gemma4, llama]

model_unfolder/adapters/transformer/families/deepseek.py

@@ -0,0 +1,107 @@
+"""Adapter for DeepSeek-V3 and Kimi K2 (which shares the architecture)."""
+from __future__ import annotations
+
+from typing import Any
+
+from ....ir import AttentionSpec, FFNSpec, ModelIR
+from ..assembly import decoder_extras, decoder_layer
+from ..common import architecture_name, get_config_value as _g, model_name
+
+
+_FAMILIES = {"deepseek_v3", "deepseek_v2", "kimi"}
+
+
+def matches(cfg: Any) -> bool:
+    arches = _g(cfg, "architectures") or []
+    model_type = _g(cfg, "model_type", "")
+    if any("DeepseekV3" in arch or "DeepseekV2" in arch or "Kimi" in arch for arch in arches):
+        return True
+    if model_type in _FAMILIES:
+        return True
+    return False
+
+
+def parse(cfg: Any) -> ModelIR:
+    num_layers = _g(cfg, "num_hidden_layers", 0)
+    hidden_size = _g(cfg, "hidden_size", 0)
+    num_heads = _g(cfg, "num_attention_heads", 0)
+    num_kv_heads = _g(cfg, "num_key_value_heads", num_heads)
+
+    kv_lora_rank = _g(cfg, "kv_lora_rank")
+    q_lora_rank = _g(cfg, "q_lora_rank")
+    qk_rope_head_dim = _g(cfg, "qk_rope_head_dim")
+    qk_nope_head_dim = _g(cfg, "qk_nope_head_dim")
+    v_head_dim = _g(cfg, "v_head_dim")
+
+    attn_kind = "mla" if kv_lora_rank is not None else "gqa"
+    head_dim = (qk_nope_head_dim or 0) + (qk_rope_head_dim or 0) or _g(cfg, "head_dim")
+
+    intermediate_size = _g(cfg, "intermediate_size", 0)
+    moe_intermediate_size = _g(cfg, "moe_intermediate_size", intermediate_size)
+    n_routed_experts = _g(cfg, "n_routed_experts") or _g(cfg, "num_experts", 0) or 0
+    n_shared_experts = _g(cfg, "n_shared_experts") or _g(cfg, "num_shared_experts", 0) or 0
+    num_experts_per_tok = _g(cfg, "num_experts_per_tok", 0)
+    first_k_dense_replace = _g(cfg, "first_k_dense_replace", 0)
+    moe_layer_freq = _g(cfg, "moe_layer_freq", 1)
+    activation = (_g(cfg, "hidden_act", "silu") or "silu").lower()
+
+    arch_name = architecture_name(cfg, "deepseek_v3")
+
+    layers = []
+    for i in range(num_layers):
+        attn = AttentionSpec(
+            kind=attn_kind,
+            num_heads=num_heads,
+            num_kv_heads=num_kv_heads,
+            head_dim=head_dim,
+            kv_lora_rank=kv_lora_rank,
+            q_lora_rank=q_lora_rank,
+            rope_dim=qk_rope_head_dim,
+            mask="causal",
+        )
+        is_moe_layer = (
+            n_routed_experts > 0
+            and i >= first_k_dense_replace
+            and (i - first_k_dense_replace) % max(moe_layer_freq, 1) == 0
+        )
+        if is_moe_layer:
+            ffn = FFNSpec(
+                kind="moe",
+                activation=activation,
+                intermediate_size=moe_intermediate_size,
+                gated=True,
+                num_experts=n_routed_experts,
+                num_experts_per_tok=num_experts_per_tok,
+                num_shared_experts=n_shared_experts,
+                expert_intermediate_size=moe_intermediate_size,
+            )
+        else:
+            ffn = FFNSpec(
+                kind="dense",
+                activation=activation,
+                intermediate_size=intermediate_size,
+                gated=True,
+            )
+        layers.append(decoder_layer(i, attn, ffn, hidden_size))
+
+    vocab_size = _g(cfg, "vocab_size", 0)
+    tie_word_embeddings = bool(_g(cfg, "tie_word_embeddings", False))
+    return ModelIR(
+        name=model_name(cfg, arch_name),
+        architecture=arch_name,
+        vocab_size=vocab_size,
+        hidden_size=hidden_size,
+        max_position_embeddings=_g(cfg, "max_position_embeddings"),
+        tie_word_embeddings=tie_word_embeddings,
+        layers=layers,
+        extras=decoder_extras(
+            vocab_size,
+            hidden_size,
+            tie_word_embeddings,
+            {
+                "v_head_dim": v_head_dim,
+                "first_k_dense_replace": first_k_dense_replace,
+                "moe_layer_freq": moe_layer_freq,
+            },
+        ),
+    )