model-unfolder 0.2.0__py3-none-any.whl

model_unfolder/ir.py

@@ -0,0 +1,163 @@
+"""
+Intermediate Representation (IR) for transformer architectures.
+
+The IR is the contract between parsers (which read HuggingFace configs)
+and the renderer (which produces SVG/HTML). It is layer-aware to support
+heterogeneous architectures (Gemma sliding-window patterns, DeepSeek
+dense+MoE phase changes, YOCO/CLA cross-layer KV sharing, etc.).
+"""
+from __future__ import annotations
+from dataclasses import dataclass, field
+from typing import Any, Optional
+
+
+@dataclass
+class AttentionSpec:
+    """Specification of an attention block within a layer."""
+    kind: str                       # "mha" | "gqa" | "mqa" | "mla"
+    num_heads: int
+    num_kv_heads: Optional[int] = None
+    head_dim: Optional[int] = None
+    kv_lora_rank: Optional[int] = None
+    q_lora_rank: Optional[int] = None
+    rope_dim: Optional[int] = None
+    mask: str = "causal"            # "causal" | "sliding" | "chunked" | "global"
+    window_size: Optional[int] = None
+    kv_source_layer: Optional[int] = None   # for cross-layer KV sharing
+
+
+@dataclass
+class FFNSpec:
+    """Specification of the feed-forward block within a layer."""
+    kind: str                       # "dense" | "moe"
+    activation: str                 # "silu" | "gelu" | "relu" | "geglu" | "swiglu"
+    intermediate_size: int
+    gated: bool = True              # SwiGLU/GeGLU style gated MLP
+    num_experts: Optional[int] = None
+    num_experts_per_tok: Optional[int] = None
+    num_shared_experts: int = 0
+    expert_intermediate_size: Optional[int] = None
+
+
+@dataclass
+class LayerSpec:
+    """One transformer layer. Instances may differ across the stack."""
+    index: int
+    attention: AttentionSpec
+    ffn: FFNSpec
+    norm_kind: str = "rmsnorm"      # "rmsnorm" | "layernorm"
+    norm_placement: str = "pre"     # "pre" | "post" | "double"
+    blocks: list = field(default_factory=list)
+
+    def signature(self) -> tuple:
+        """Hashable structural fingerprint used for grouping similar layers."""
+        a = self.attention
+        f = self.ffn
+        return (
+            a.kind, a.mask, a.window_size, a.kv_source_layer is not None,
+            f.kind, f.gated, f.num_experts,
+            self.norm_kind, self.norm_placement,
+        )
+
+
+@dataclass
+class CrossLayerEdge:
+    """A dependency between two layers (e.g. KV cache sharing)."""
+    kind: str                       # "kv_share"
+    from_layer: int
+    to_layer: int
+    shared: list = field(default_factory=list)    # ["K", "V"]
+
+
+@dataclass
+class ModelIR:
+    """Top-level IR for a complete model."""
+    name: str
+    architecture: str               # e.g. "DeepseekV3ForCausalLM"
+    vocab_size: int
+    hidden_size: int
+    max_position_embeddings: Optional[int]
+    tie_word_embeddings: bool
+    layers: list                    # list[LayerSpec]
+    cross_layer_edges: list = field(default_factory=list)
+    extras: dict = field(default_factory=dict)
+
+    def to_dict(self) -> dict:
+        # Avoid dataclasses.asdict here: it recursively deepcopy()s every
+        # nested dict/list, including repeated render block metadata for every
+        # layer.  The IR is treated as immutable after parsing, so a direct
+        # structural projection is much cheaper and enough for rendering.
+        return {
+            "name": self.name,
+            "architecture": self.architecture,
+            "vocab_size": self.vocab_size,
+            "hidden_size": self.hidden_size,
+            "max_position_embeddings": self.max_position_embeddings,
+            "tie_word_embeddings": self.tie_word_embeddings,
+            "layers": [_layer_to_dict(layer) for layer in self.layers],
+            "cross_layer_edges": [_cross_edge_to_dict(edge) for edge in self.cross_layer_edges],
+            "extras": self.extras,
+        }
+
+    @property
+    def num_layers(self) -> int:
+        return len(self.layers)
+
+    def layer_groups(self) -> list:
+        """Run-length encode layers by signature."""
+        groups = []
+        for layer in self.layers:
+            sig = layer.signature()
+            if groups and groups[-1][0] == sig:
+                groups[-1][1].append(layer.index)
+            else:
+                groups.append((sig, [layer.index]))
+        return groups
+
+
+def _attention_to_dict(a: AttentionSpec) -> dict:
+    return {
+        "kind": a.kind,
+        "num_heads": a.num_heads,
+        "num_kv_heads": a.num_kv_heads,
+        "head_dim": a.head_dim,
+        "kv_lora_rank": a.kv_lora_rank,
+        "q_lora_rank": a.q_lora_rank,
+        "rope_dim": a.rope_dim,
+        "mask": a.mask,
+        "window_size": a.window_size,
+        "kv_source_layer": a.kv_source_layer,
+    }
+
+
+def _ffn_to_dict(f: FFNSpec) -> dict:
+    return {
+        "kind": f.kind,
+        "activation": f.activation,
+        "intermediate_size": f.intermediate_size,
+        "gated": f.gated,
+        "num_experts": f.num_experts,
+        "num_experts_per_tok": f.num_experts_per_tok,
+        "num_shared_experts": f.num_shared_experts,
+        "expert_intermediate_size": f.expert_intermediate_size,
+    }
+
+
+def _layer_to_dict(layer: LayerSpec) -> dict:
+    return {
+        "index": layer.index,
+        "attention": _attention_to_dict(layer.attention),
+        "ffn": _ffn_to_dict(layer.ffn),
+        "norm_kind": layer.norm_kind,
+        "norm_placement": layer.norm_placement,
+        "blocks": layer.blocks,
+    }
+
+
+def _cross_edge_to_dict(edge: CrossLayerEdge) -> dict:
+    return {
+        "kind": edge.kind,
+        "from_layer": edge.from_layer,
+        "to_layer": edge.to_layer,
+        "shared": edge.shared,
+    }

model_unfolder/labels.py

@@ -0,0 +1,166 @@
+"""Renderer-agnostic vocabulary for talking about transformer specs.
+
+Anything that needs to refer to attention masks ("SWA", "full"), attention
+kinds ("GQA"), or build a human description of an attention / FFN block goes
+through this module.  Keeping it in one place means changing the wording
+(swap "SWA" for "Sliding", say) only requires editing this file — no scavenger
+hunt across the renderer.
+
+The functions are pure and operate on plain ``dict``-shaped specs (i.e. what
+``ModelIR.to_dict()`` produces), so they're equally useful inside the HTML
+renderer, a future text/markdown renderer, or notebook utilities.
+"""
+from __future__ import annotations
+
+_MASK_SHORT = {
+    "sliding": "SWA",
+    "global": "full",
+    "causal": "causal",
+    "chunked": "chunked",
+}
+_MASK_LONG = {
+    "sliding": "Sliding-window",
+    "global": "Full / global",
+    "causal": "Causal",
+    "chunked": "Chunked",
+}
+_MASK_TITLE = {
+    "sliding": "Sliding-window attention",
+    "global": "Full-context attention",
+    "causal": "Causal attention",
+    "chunked": "Chunked attention",
+}
+
+_KIND_SHORT = {"mla": "MLA", "gqa": "GQA", "mqa": "MQA", "mha": "MHA"}
+_KIND_LONG = {
+    "mla": "Multi-head latent attention",
+    "gqa": "Grouped-query attention",
+    "mqa": "Multi-query attention",
+    "mha": "Multi-head attention",
+}
+_ACTIVATION_LABELS = {
+    "gelu": "GELU",
+    "gelu_new": "GELU",
+    "gelu_fast": "GELU",
+    "gelu_pytorch_tanh": "GELU",
+    "relu": "ReLU",
+    "silu": "SiLU",
+    "swish": "SiLU",
+    "geglu": "GEGLU",
+    "swiglu": "SwiGLU",
+}
+
+
+def mask_short(attention: dict) -> str:
+    """Compact mask tag — ``"SWA"`` / ``"full"`` / ``"causal"``."""
+    return _MASK_SHORT.get(attention.get("mask", "causal"), "causal")
+
+
+def mask_long(attention: dict) -> str:
+    """Human-readable mask label — ``"Sliding-window"`` / ``"Full / global"``."""
+    return _MASK_LONG.get(attention.get("mask", "causal"), "Causal")
+
+
+def mask_title(attention: dict) -> str:
+    """Tooltip-style mask description."""
+    return _MASK_TITLE.get(attention.get("mask", "causal"), "Causal attention")
+
+
+def mask_chip(attention: dict) -> str:
+    """One-line chip; for sliding includes the window size: ``"SWA 1,024"``."""
+    label = mask_short(attention)
+    window = attention.get("window_size")
+    if window and is_sliding(attention):
+        return f"{label} {_fmt_int(window)}"
+    return label
+
+
+def kind_short(attention: dict) -> str:
+    return _KIND_SHORT.get(attention.get("kind", ""), "MHA")
+
+
+def kind_long(attention: dict) -> str:
+    return _KIND_LONG.get(attention.get("kind", ""), "Multi-head attention")
+
+
+def is_sliding(attention: dict) -> bool:
+    return attention.get("mask") == "sliding"
+
+
+def is_global(attention: dict) -> bool:
+    return attention.get("mask") == "global"
+
+
+def kv_shared(attention: dict) -> bool:
+    """True when this layer reuses K/V from an earlier layer (Gemma 4 small)."""
+    return attention.get("kv_source_layer") is not None
+
+
+def activation_label(name: str | None) -> str:
+    """Display label for activation names stored in configs.
+
+    Configs often expose backend-specific names such as
+    ``gelu_pytorch_tanh``.  Diagrams should show the mathematical operation,
+    not the implementation detail.
+    """
+    key = (name or "").strip().lower().replace("-", "_")
+    if key in _ACTIVATION_LABELS:
+        return _ACTIVATION_LABELS[key]
+    if key.startswith("gelu"):
+        return "GELU"
+    return key.replace("_", " ").title() if key else "Activation"
+
+
+def describe_attention(attention: dict) -> str:
+    """Multi-clause human description suitable for tooltips and cards.
+
+    The window-size suffix is appended whenever the layer uses sliding-window
+    attention, regardless of the underlying kind (GQA/MHA/MLA).
+    """
+    kind = attention.get("kind")
+    if kind == "mla":
+        text = (
+            f"Multi-head latent attention; {attention.get('num_heads')} heads; "
+            f"KV LoRA {_fmt_int(attention.get('kv_lora_rank'))}"
+        )
+        if attention.get("q_lora_rank"):
+            text += f"; Q LoRA {_fmt_int(attention.get('q_lora_rank'))}"
+    elif kind == "gqa":
+        text = (
+            f"Grouped-query; {attention.get('num_heads')} Q / "
+            f"{attention.get('num_kv_heads')} KV heads; "
+            f"head dim {_fmt_int(attention.get('head_dim'))}"
+        )
+    elif kind == "mqa":
+        text = f"Multi-query; {attention.get('num_heads')} Q / 1 KV head"
+    else:
+        text = (
+            f"Multi-head; {attention.get('num_heads')} heads; "
+            f"head dim {_fmt_int(attention.get('head_dim'))}"
+        )
+    if is_sliding(attention) and attention.get("window_size"):
+        text += f"; sliding window {_fmt_int(attention.get('window_size'))}"
+    return text
+
+
+def describe_ffn(ffn: dict) -> str:
+    """Multi-clause human description of an FFN / MoE block."""
+    if ffn.get("kind") == "moe":
+        text = f"MoE; {_fmt_int(ffn.get('num_experts'))} experts; top-{ffn.get('num_experts_per_tok')}"
+        if ffn.get("num_shared_experts"):
+            text += f" + {ffn.get('num_shared_experts')} shared"
+        if ffn.get("num_experts") and ffn.get("num_experts_per_tok"):
+            text += f"; {100 * ffn['num_experts_per_tok'] / ffn['num_experts']:.1f}% active"
+        text += f"; expert hidden {_fmt_int(ffn.get('expert_intermediate_size') or ffn.get('intermediate_size'))}"
+        return text
+    gated = "gated " if ffn.get("gated") else ""
+    return f"{gated}FFN; {activation_label(ffn.get('activation'))}; hidden {_fmt_int(ffn.get('intermediate_size'))}"
+
+
+def _fmt_int(value) -> str:
+    if value is None:
+        return "?"
+    try:
+        return f"{int(value):,}"
+    except (TypeError, ValueError):
+        return str(value)

model_unfolder/params.py

@@ -0,0 +1,119 @@
+"""Rough parameter-count estimation from an IR.
+
+These are *estimates*. We don't try to model every implementation detail
+(bias terms, MLA's exact projection layout, expert grouping, etc.) — just
+enough to give the right order of magnitude and a useful active/total split
+for MoE models.
+"""
+from __future__ import annotations
+from .ir import ModelIR, AttentionSpec, FFNSpec
+
+
+def _attn_params(a: AttentionSpec, hidden: int) -> int:
+    h = hidden
+    head_dim = a.head_dim or (h // max(a.num_heads, 1))
+    nq = a.num_heads
+    nkv = a.num_kv_heads or nq
+
+    if a.kind == "mla":
+        # Q path: optional LoRA down then up to (nope+rope)*nq
+        q_out = nq * head_dim
+        if a.q_lora_rank:
+            q = h * a.q_lora_rank + a.q_lora_rank * q_out
+        else:
+            q = h * q_out
+        # KV path: hidden -> (kv_lora_rank + rope_dim), then up to nq*(nope+v)
+        kv_lora = a.kv_lora_rank or 0
+        rope = a.rope_dim or 0
+        kv_down = h * (kv_lora + rope)
+        kv_up = kv_lora * (nq * head_dim * 2)  # K nope + V — rough
+        kv = kv_down + kv_up
+        out = nq * head_dim * h
+        return q + kv + out
+
+    qkv = h * (nq + 2 * nkv) * head_dim
+    out = nq * head_dim * h
+    return qkv + out
+
+
+def _ffn_params(f: FFNSpec, hidden: int) -> tuple:
+    """Returns (total_params, active_params_per_token)."""
+    g = 3 if f.gated else 2
+    if f.kind == "moe":
+        per_expert = g * hidden * (f.expert_intermediate_size or f.intermediate_size)
+        n_routed = f.num_experts or 0
+        n_shared = f.num_shared_experts or 0
+        n_active = f.num_experts_per_tok or 0
+        router = hidden * n_routed
+        total = per_expert * (n_routed + n_shared) + router
+        active = per_expert * (n_active + n_shared) + router
+        return total, active
+    p = g * hidden * f.intermediate_size
+    return p, p
+
+
+def estimate_params(ir: ModelIR) -> dict:
+    """Estimate parameter counts for a model.
+
+    Returns a dict::
+
+        {
+            "total":     int,   # all parameters
+            "active":    int,   # active per token (== total for non-MoE)
+            "embed":     int,
+            "output":    int,
+            "per_layer": [{"total": int, "active": int}, ...],
+            "is_sparse": bool,
+        }
+    """
+    h = ir.hidden_size
+    v = ir.vocab_size
+
+    embed = v * h
+    output = 0 if ir.tie_word_embeddings else v * h
+    final_norm = h
+
+    per_layer = []
+    layers_total = 0
+    layers_active = 0
+    is_sparse = False
+
+    for layer in ir.layers:
+        a_p = _attn_params(layer.attention, h)
+        f_total, f_active = _ffn_params(layer.ffn, h)
+        if layer.ffn.kind == "moe":
+            is_sparse = True
+        norm_p = 2 * h
+        t = a_p + f_total + norm_p
+        ac = a_p + f_active + norm_p
+        per_layer.append({"total": t, "active": ac, "attn": a_p, "ffn": f_total})
+        layers_total += t
+        layers_active += ac
+
+    total = embed + output + final_norm + layers_total
+    active = embed + output + final_norm + layers_active
+
+    return {
+        "total": total,
+        "active": active,
+        "embed": embed,
+        "output": output,
+        "per_layer": per_layer,
+        "is_sparse": is_sparse,
+    }
+
+
+def humanize(n: int) -> str:
+    """Format a parameter count as e.g. '671B', '37.4B', '8.2M'."""
+    if n is None:
+        return "?"
+    n = float(n)
+    for unit, scale in (("T", 1e12), ("B", 1e9), ("M", 1e6), ("K", 1e3)):
+        if n >= scale:
+            v = n / scale
+            if v >= 100:
+                return f"{v:.0f}{unit}"
+            if v >= 10:
+                return f"{v:.1f}{unit}"
+            return f"{v:.2f}{unit}"
+    return f"{int(n)}"

model_unfolder/parser.py

@@ -0,0 +1,137 @@
+"""Parse a HuggingFace config (or model ID) into our IR."""
+from __future__ import annotations
+import os
+from typing import Any
+from .ir import ModelIR
+from .adapters import find_adapter
+
+
+HF_TOKEN_ENV_VARS = (
+    "HF_TOKEN",
+    "HUGGING_FACE_HUB_TOKEN",
+    "HUGGINGFACE_HUB_TOKEN",
+)
+
+
+def config_to_ir(cfg_or_id: Any, token: Any = None) -> ModelIR:
+    """Parse anything HF-shaped into an IR.
+
+    Accepts:
+        - A HuggingFace ``PretrainedConfig`` instance
+        - A model ID string (e.g. ``"moonshotai/Kimi-K2-Instruct"``) — requires ``transformers``
+        - A plain ``dict`` (the contents of ``config.json``)
+
+    Parameters
+    ----------
+    token
+        Optional Hugging Face token used only when loading a config by model ID.
+        If omitted, ``HF_TOKEN`` and legacy Hugging Face token env vars are used
+        when present.
+    """
+    cfg = _coerce(cfg_or_id, token=token)
+    adapter = find_adapter(cfg)
+    if adapter is None:
+        arches = (
+            cfg.get("architectures") if isinstance(cfg, dict)
+            else getattr(cfg, "architectures", None)
+        )
+        raise ValueError(
+            f"No adapter found for architecture {arches}. "
+            "Pass a dict-like config or contribute an adapter."
+        )
+    return adapter.parse(cfg)
+
+
+def _coerce(cfg_or_id, token: Any = None):
+    if isinstance(cfg_or_id, dict):
+        return cfg_or_id
+    if isinstance(cfg_or_id, str):
+        try:
+            from transformers import AutoConfig
+        except ImportError as e:
+            raise ImportError(
+                "Loading a model by ID requires `transformers`. "
+                "Install with `pip install transformers`, or pass a config dict."
+            ) from e
+        return _load_config_from_hf(AutoConfig, cfg_or_id, token=token)
+    return cfg_or_id
+
+
+def _load_config_from_hf(auto_config: Any, model_id: str, token: Any = None):
+    auth_token = _resolve_hf_token(token)
+    try:
+        return _from_pretrained(auto_config, model_id, auth_token, trust_remote_code=False)
+    except Exception as e:
+        if not _should_retry_with_remote_code(e):
+            raise
+        return _from_pretrained(auto_config, model_id, auth_token, trust_remote_code=True)
+
+
+def _from_pretrained(auto_config: Any, model_id: str, auth_token: Any, *, trust_remote_code: bool):
+    kwargs = {}
+    if trust_remote_code:
+        kwargs["trust_remote_code"] = True
+    if auth_token is None:
+        return auto_config.from_pretrained(model_id, **kwargs)
+
+    try:
+        return auto_config.from_pretrained(model_id, token=auth_token, **kwargs)
+    except Exception as e:
+        if not _should_retry_with_legacy_auth(e):
+            raise
+        return auto_config.from_pretrained(
+            model_id,
+            use_auth_token=auth_token,
+            **kwargs,
+        )
+
+
+def _resolve_hf_token(token: Any = None):
+    if token is not None:
+        return _clean_token(token)
+    for name in HF_TOKEN_ENV_VARS:
+        value = _clean_token(os.environ.get(name))
+        if value is not None:
+            return value
+    return None
+
+
+def _clean_token(token: Any):
+    if isinstance(token, str):
+        token = token.strip()
+        return token or None
+    return token
+
+
+def _should_retry_with_legacy_auth(error: Exception) -> bool:
+    msg = str(error).lower()
+    return any(
+        marker in msg
+        for marker in (
+            "token",
+            "use_auth_token",
+            "authentication",
+            "authorization",
+            "unauthorized",
+            "forbidden",
+            "401",
+            "403",
+            "gated",
+            "private",
+        )
+    )
+
+
+def _should_retry_with_remote_code(error: Exception) -> bool:
+    msg = str(error).lower()
+    return any(
+        marker in msg
+        for marker in (
+            "trust_remote_code",
+            "remote code",
+            "custom code",
+            "custom configuration",
+            "execute the configuration file",
+            "execute the repository",
+        )
+    )

model_unfolder/renderers/__init__.py

@@ -0,0 +1 @@
+"""Rendering backends for model-unfolder diagrams."""

model_unfolder/renderers/html/__init__.py

@@ -0,0 +1,5 @@
+"""HTML/SVG rendering backend."""
+
+from .document import render_document, render_fragment
+
+__all__ = ["render_document", "render_fragment"]

model_unfolder/renderers/html/block_views/__init__.py

@@ -0,0 +1,20 @@
+"""Reusable rich block detail views for the HTML renderer."""
+from __future__ import annotations
+
+from .attention import attention_card, attention_card_css
+from .feed_forward import build_ffn_view, build_moe_view
+from .per_layer_embedding import build_per_layer_embedding_view
+
+
+def block_detail_svg(ir: dict, info: dict, mount_id: str, block: dict) -> str | None:
+    """Return a rich SVG for a clicked architecture block, when one exists."""
+    if block.get("kind") == "ffn":
+        ffn = info["dominant"]["spec"]["ffn"]
+        if ffn.get("kind") == "moe":
+            return build_moe_view(ir, info, mount_id)
+        return build_ffn_view(ir, info, mount_id)
+
+    if block.get("detail_view") == "per_layer_embedding":
+        return build_per_layer_embedding_view(ir, info, mount_id, block)
+
+    return None