model-unfolder 0.2.0__py3-none-any.whl

model_unfolder/adapters/transformer/families/gemma4.py

@@ -0,0 +1,202 @@
+"""Adapter for Google's Gemma 4 family (E2B / E4B / 31B / 26B-A4B).
+
+Gemma 4 nests its language-model fields under ``text_config`` since the
+top-level config also covers the vision and audio encoders.  The text stack
+itself has a few features the older Llama/Gemma adapters don't model:
+
+* Per-layer ``layer_types`` array tagging each block as ``sliding_attention``
+  or ``full_attention`` (instead of Gemma 3's "every Nth layer" pattern).
+* Dual attention shape: sliding layers use ``head_dim`` + ``num_key_value_heads``;
+  global layers use ``global_head_dim`` + ``num_global_key_value_heads``.
+* Optional MoE FFN (26B-A4B) controlled by ``enable_moe_block``.
+* Per-Layer Embeddings (PLE) on small models — ``hidden_size_per_layer_input``
+  is the parallel conditioning dim, ``vocab_size_per_layer_input`` its vocab.
+* Shared-KV layers — the last ``num_kv_shared_layers`` layers reuse K/V from
+  the last earlier layer of the same attention type.
+"""
+from __future__ import annotations
+
+from typing import Any
+
+from ....ir import AttentionSpec, CrossLayerEdge, FFNSpec, ModelIR
+from ..assembly import decoder_extras, decoder_layer
+from ..common import architecture_name, get_config_value as _g, model_name
+from ..special_parts.per_layer_embedding import (
+    per_layer_embedding_blocks,
+    per_layer_embedding_extras,
+)
+
+
+_TOP_TYPES = {"gemma4"}
+_TEXT_TYPES = {"gemma4_text"}
+_ARCH_HINTS = ("gemma4",)
+
+
+def matches(cfg: Any) -> bool:
+    arches = _g(cfg, "architectures") or []
+    if any(any(hint in arch.lower() for hint in _ARCH_HINTS) for arch in arches):
+        return True
+    model_type = (_g(cfg, "model_type", "") or "").lower()
+    if model_type in _TOP_TYPES or model_type in _TEXT_TYPES:
+        return True
+    return False
+
+
+def parse(cfg: Any) -> ModelIR:
+    arch_name = architecture_name(cfg, "gemma4")
+    text_cfg = _text_config(cfg)
+
+    num_layers = _g(text_cfg, "num_hidden_layers", 0)
+    hidden_size = _g(text_cfg, "hidden_size", 0)
+    intermediate_size = _g(text_cfg, "intermediate_size", 0)
+    activation = (_g(text_cfg, "hidden_activation") or _g(text_cfg, "hidden_act") or "gelu_pytorch_tanh").lower()
+
+    num_q = _g(text_cfg, "num_attention_heads", 0)
+    num_kv = _g(text_cfg, "num_key_value_heads", num_q)
+    num_kv_global = _g(text_cfg, "num_global_key_value_heads") or num_kv
+    head_dim = _g(text_cfg, "head_dim") or (hidden_size // num_q if num_q else None)
+    head_dim_global = _g(text_cfg, "global_head_dim") or head_dim
+    sliding_window = _g(text_cfg, "sliding_window")
+
+    layer_types = _g(text_cfg, "layer_types") or []
+
+    moe_enabled = bool(_g(text_cfg, "enable_moe_block"))
+    num_experts = _g(text_cfg, "num_experts") or 0
+    top_k = _g(text_cfg, "top_k_experts") or 0
+    moe_intermediate_size = _g(text_cfg, "moe_intermediate_size") or 0
+
+    num_kv_shared = _g(text_cfg, "num_kv_shared_layers") or 0
+    first_shared = num_layers - num_kv_shared if num_kv_shared else num_layers
+
+    ple_dim = _g(text_cfg, "hidden_size_per_layer_input") or 0
+    ple_vocab = _g(text_cfg, "vocab_size_per_layer_input") or _g(text_cfg, "vocab_size", 0)
+
+    layers = []
+    cross_edges: list[CrossLayerEdge] = []
+    for i in range(num_layers):
+        layer_type = layer_types[i] if i < len(layer_types) else "full_attention"
+        is_sliding = "sliding" in layer_type
+
+        if is_sliding:
+            mask = "sliding"
+            window = sliding_window
+            kv_heads = num_kv
+            this_head_dim = head_dim
+        else:
+            mask = "global"
+            window = None
+            kv_heads = num_kv_global
+            this_head_dim = head_dim_global
+
+        kv_source: int | None = None
+        if i >= first_shared:
+            kv_source = _last_matching_layer(layer_types, i, first_shared)
+            if kv_source is not None:
+                cross_edges.append(
+                    CrossLayerEdge(
+                        kind="kv_share",
+                        from_layer=kv_source,
+                        to_layer=i,
+                        shared=["K", "V"],
+                    )
+                )
+
+        attn_kind = _attention_kind(num_q, kv_heads)
+        attn = AttentionSpec(
+            kind=attn_kind,
+            num_heads=num_q,
+            num_kv_heads=kv_heads,
+            head_dim=this_head_dim,
+            mask=mask,
+            window_size=window,
+            kv_source_layer=kv_source,
+        )
+
+        if moe_enabled and num_experts:
+            ffn = FFNSpec(
+                kind="moe",
+                activation=activation,
+                intermediate_size=intermediate_size,
+                gated=True,
+                num_experts=num_experts,
+                num_experts_per_tok=top_k,
+                expert_intermediate_size=moe_intermediate_size or intermediate_size,
+            )
+        else:
+            ffn = FFNSpec(
+                kind="dense",
+                activation=activation,
+                intermediate_size=intermediate_size,
+                gated=True,
+            )
+
+        extra_blocks = []
+        if ple_dim:
+            extra_blocks.extend(
+                per_layer_embedding_blocks(hidden_size, ple_dim, activation="gelu")
+            )
+        layers.append(
+            decoder_layer(i, attn, ffn, hidden_size, extra_blocks=extra_blocks)
+        )
+
+    vocab_size = _g(text_cfg, "vocab_size", 0)
+    tie_word_embeddings = bool(_g(text_cfg, "tie_word_embeddings", _g(cfg, "tie_word_embeddings", False)))
+
+    extras = decoder_extras(
+        vocab_size,
+        hidden_size,
+        tie_word_embeddings,
+        per_layer_embedding_extras(hidden_size, ple_dim, ple_vocab, num_layers)
+        if ple_dim else None,
+    )
+    if num_kv_shared:
+        extras["num_kv_shared_layers"] = num_kv_shared
+    if _g(text_cfg, "attention_k_eq_v"):
+        extras["attention_k_eq_v"] = True
+    if _g(text_cfg, "use_double_wide_mlp"):
+        extras["use_double_wide_mlp"] = True
+
+    return ModelIR(
+        name=model_name(cfg, arch_name),
+        architecture=arch_name,
+        vocab_size=vocab_size,
+        hidden_size=hidden_size,
+        max_position_embeddings=_g(text_cfg, "max_position_embeddings"),
+        tie_word_embeddings=tie_word_embeddings,
+        layers=layers,
+        cross_layer_edges=cross_edges,
+        extras=extras,
+    )
+
+
+def _text_config(cfg: Any) -> Any:
+    """Reach the language-model sub-config when ``cfg`` is the multimodal wrapper."""
+    if (_g(cfg, "model_type", "") or "").lower() in _TEXT_TYPES:
+        return cfg
+    sub = _g(cfg, "text_config")
+    return sub if sub is not None else cfg
+
+
+def _attention_kind(num_q: int, num_kv: int) -> str:
+    if not num_q:
+        return "mha"
+    if num_kv == num_q:
+        return "mha"
+    if num_kv == 1:
+        return "mqa"
+    return "gqa"
+
+
+def _last_matching_layer(layer_types: list, i: int, first_shared: int) -> int | None:
+    """Source layer for a KV-shared layer.
+
+    Per the Gemma 4 release notes: shared layers reuse K/V from the most
+    recent non-shared layer of the *same* attention type (sliding or full).
+    """
+    if not layer_types or i >= len(layer_types):
+        return None
+    target_type = layer_types[i]
+    for j in range(min(first_shared, len(layer_types)) - 1, -1, -1):
+        if layer_types[j] == target_type:
+            return j
+    return None

model_unfolder/adapters/transformer/families/llama.py

@@ -0,0 +1,91 @@
+"""Adapter for Llama, Mistral, Qwen, and similar GQA/MHA dense models."""
+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 = {"llama", "mistral", "qwen2", "qwen3", "phi3", "gemma"}
+
+
+def matches(cfg: Any) -> bool:
+    arches = _g(cfg, "architectures") or []
+    model_type = _g(cfg, "model_type", "")
+    for arch in arches:
+        if any(fam in arch.lower() for fam in ("llama", "mistral", "qwen", "phi3")):
+            return True
+    if model_type in _FAMILIES:
+        return True
+    return False
+
+
+def parse(cfg: Any) -> ModelIR:
+    num_layers = _g(cfg, "num_hidden_layers", 0)
+    num_heads = _g(cfg, "num_attention_heads", 0)
+    num_kv_heads = _g(cfg, "num_key_value_heads", num_heads)
+    hidden_size = _g(cfg, "hidden_size", 0)
+    head_dim = _g(cfg, "head_dim") or (hidden_size // num_heads if num_heads else None)
+
+    if num_kv_heads == num_heads:
+        attn_kind = "mha"
+    elif num_kv_heads == 1:
+        attn_kind = "mqa"
+    else:
+        attn_kind = "gqa"
+
+    sliding_window = _g(cfg, "sliding_window")
+    sliding_pattern = _g(cfg, "sliding_window_pattern")
+    layer_types = _g(cfg, "layer_types")
+
+    intermediate_size = _g(cfg, "intermediate_size", 0)
+    activation = (_g(cfg, "hidden_act", "silu") or "silu").lower()
+
+    arch_name = architecture_name(cfg, "llama")
+
+    layers = []
+    for i in range(num_layers):
+        if layer_types and i < len(layer_types):
+            layer_type = layer_types[i]
+            if "sliding" in layer_type:
+                mask, win = "sliding", sliding_window
+            else:
+                mask, win = "causal", None
+        elif sliding_pattern and sliding_window:
+            mask = "sliding" if (i % sliding_pattern) != (sliding_pattern - 1) else "causal"
+            win = sliding_window if mask == "sliding" else None
+        elif sliding_window:
+            mask, win = "sliding", sliding_window
+        else:
+            mask, win = "causal", None
+
+        attn = AttentionSpec(
+            kind=attn_kind,
+            num_heads=num_heads,
+            num_kv_heads=num_kv_heads,
+            head_dim=head_dim,
+            mask=mask,
+            window_size=win,
+        )
+        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),
+    )

model_unfolder/adapters/transformer/special_parts/__init__.py

@@ -0,0 +1,2 @@
+"""Reusable transformer layer parts."""
+

model_unfolder/adapters/transformer/special_parts/per_layer_embedding.py

@@ -0,0 +1,220 @@
+"""Reusable Per-Layer Embedding (PLE) transformer part.
+
+This module owns the canonical IR shape for PLE-style conditioning.  A model
+family adapter should only detect that the config has such a pathway, then
+attach these blocks and extras.  The renderer consumes the declared block
+metadata, so the feature is not tied to any one model family.
+"""
+from __future__ import annotations
+
+from ....labels import activation_label
+from ..common import format_dim as _fmt
+
+
+DEFAULT_BLOCK_ID = "ple"
+DEFAULT_ADD_ID = "add3"
+DEFAULT_PATHWAY_ID = "per_layer_input"
+
+
+def per_layer_embedding_blocks(
+    hidden_size: int,
+    embedding_dim: int,
+    *,
+    activation: str = "gelu",
+    block_id: str = DEFAULT_BLOCK_ID,
+    add_id: str = DEFAULT_ADD_ID,
+    pathway_id: str = DEFAULT_PATHWAY_ID,
+    lane: str = "left",
+    tap_from: str = "rms1",
+    feeds: str | None = None,
+    residual_from: str = "add2",
+) -> list[dict]:
+    """Return the layer blocks for a reusable PLE side pathway.
+
+    The canonical shape is:
+
+    hidden -> gate -> activation -> multiply(per-layer vector) -> projection
+    -> norm -> residual add.
+
+    The side block is intentionally rendered off the central chain via
+    ``lane``/``tap_from``/``feeds``.  Future adapters can reuse this exact
+    shape without adding renderer branches for their model family.
+    """
+    feeds = feeds or add_id
+    hidden = _fmt(hidden_size)
+    emb = _fmt(embedding_dim)
+    act_name = _activation_label(activation)
+    ids = _child_ids(block_id)
+
+    children = [
+        {
+            "id": ids["gate"],
+            "label": "Linear (gate)",
+            "title": "Per-layer input gate",
+            "description": f"Linear; {hidden} -> {emb}",
+        },
+        {
+            "id": ids["activation"],
+            "label": act_name,
+            "title": "PLE activation",
+            "description": f"Element-wise {act_name}",
+        },
+        {
+            "id": ids["multiply"],
+            "label": "x",
+            "title": "Per-layer gate (x)",
+            "description": (
+                f"Element-wise multiply by {pathway_id}[L] "
+                f"({emb}-d vector sourced from the parallel pathway)"
+            ),
+        },
+        {
+            "id": pathway_id,
+            "label": f"{pathway_id}[L]",
+            "title": "Per-layer input vector",
+            "description": f"{emb}-d vector produced outside the layer stack for layer L.",
+        },
+        {
+            "id": ids["projection"],
+            "label": "Linear (up)",
+            "title": "Per-layer projection",
+            "description": f"Linear; {emb} -> {hidden}",
+        },
+        {
+            "id": ids["norm"],
+            "label": "RMSNorm",
+            "title": "Post-PLE norm",
+            "description": f"RMSNorm; dim {hidden}",
+        },
+    ]
+
+    return [
+        {
+            "id": block_id,
+            "role": "ple",
+            "kind": "ple",
+            "label": "PLE",
+            "title": "Per-Layer Embeddings",
+            "description": (
+                f"Per-layer gate-and-project; {hidden} -> {emb} -> {hidden}. "
+                "Multiplied by a per-layer vector built outside the stack."
+            ),
+            "detail_view": "per_layer_embedding",
+            "detail": {
+                "view": "per_layer_embedding",
+                "view_id": block_id,
+                "pathway_id": pathway_id,
+                "nodes": ids,
+                "input_label": "in  (hidden)",
+                "output_label": "out  -> add (residual)",
+                "external_label": f"{pathway_id}[L]",
+                "external_description": f"({emb}-d, built outside layers)",
+                "hidden_size": hidden_size,
+                "embedding_dim": embedding_dim,
+            },
+            "lane": lane,
+            "tap_from": tap_from,
+            "feeds": feeds,
+            "children": children,
+        },
+        {
+            "id": add_id,
+            "role": "residual",
+            "kind": "residual_add",
+            "residual_from": residual_from,
+            "label": "+",
+            "title": "Residual add (PLE)",
+            "description": "post-FFN + PLE output",
+        },
+    ]
+
+
+def per_layer_embedding_pathway(
+    hidden_size: int,
+    embedding_dim: int,
+    vocab_size: int,
+    num_layers: int,
+    *,
+    pathway_id: str = DEFAULT_PATHWAY_ID,
+    block_id: str = DEFAULT_BLOCK_ID,
+) -> dict:
+    """Return the external pathway descriptor consumed by the PLE block."""
+    hidden = _fmt(hidden_size)
+    emb = _fmt(embedding_dim)
+    vocab = _fmt(vocab_size)
+    layers = _fmt(num_layers)
+    ids = _child_ids(block_id)
+    return {
+        "id": pathway_id,
+        "label": "Per-Layer Embeddings",
+        "short_label": "PLE",
+        "description": (
+            f"Parallel pathway producing one {emb}-d vector per layer per token; "
+            "feeds every layer's PLE gate."
+        ),
+        "feeds": "every_layer",
+        "tap_block": ids["multiply"],
+        "construction": [
+            {
+                "id": f"{block_id}_lookup",
+                "label": "embed_tokens_per_layer",
+                "kind": "embedding",
+                "description": f"Lookup; {vocab} -> {layers} x {emb}",
+            },
+            {
+                "id": f"{block_id}_proj_in",
+                "label": "per_layer_model_projection",
+                "kind": "linear",
+                "description": f"Linear; {hidden} -> {layers} x {emb}",
+            },
+            {
+                "id": f"{block_id}_combine",
+                "label": "(token + context) / sqrt(2)",
+                "kind": "scale_add",
+                "description": "Sum the two pathways and rescale.",
+            },
+        ],
+    }
+
+
+def per_layer_embedding_extras(
+    hidden_size: int,
+    embedding_dim: int,
+    vocab_size: int,
+    num_layers: int,
+    *,
+    pathway_id: str = DEFAULT_PATHWAY_ID,
+    block_id: str = DEFAULT_BLOCK_ID,
+) -> dict:
+    """Return top-level IR extras for a reusable PLE pathway."""
+    return {
+        "per_layer_embeddings": {
+            "hidden": embedding_dim,
+            "vocab": vocab_size,
+            "pathway_id": pathway_id,
+        },
+        "external_pathways": [
+            per_layer_embedding_pathway(
+                hidden_size,
+                embedding_dim,
+                vocab_size,
+                num_layers,
+                pathway_id=pathway_id,
+                block_id=block_id,
+            )
+        ],
+    }
+
+
+def _child_ids(block_id: str) -> dict[str, str]:
+    return {
+        "gate": f"{block_id}_gate",
+        "activation": f"{block_id}_act",
+        "multiply": f"{block_id}_mul",
+        "projection": f"{block_id}_proj",
+        "norm": f"{block_id}_norm",
+    }
+
+
+def _activation_label(activation: str) -> str:
+    return activation_label(activation or "gelu")

model_unfolder/diagram.py

@@ -0,0 +1,95 @@
+"""Diagram — the renderable object.
+
+Implements ``_repr_html_`` so it auto-renders inline in Jupyter (like
+``matplotlib`` or a ``pandas`` DataFrame). Outside notebooks, call
+``.save(path)`` to write a portable HTML file.
+"""
+from __future__ import annotations
+import json
+import os
+import uuid
+from .ir import ModelIR
+from .html_renderer import render_document, render_fragment
+from .params import estimate_params, humanize
+
+
+class Diagram:
+    """A renderable diagram of a transformer architecture."""
+
+    def __init__(self, ir: ModelIR):
+        self.ir = ir
+        self._mount_id = f"uf-{uuid.uuid4().hex[:10]}"
+        self._params = estimate_params(ir)
+        self._ir_cache: dict | None = None
+        self._html_cache: dict[bool, str] = {}
+
+    def to_ir(self) -> dict:
+        """Return the underlying IR (plus param estimates) as a plain dict."""
+        if self._ir_cache is not None:
+            return self._ir_cache
+
+        d = self.ir.to_dict()
+        p = self._params
+        d["params"] = {
+            "total": p["total"],
+            "active": p["active"],
+            "total_h": humanize(p["total"]),
+            "active_h": humanize(p["active"]),
+            "is_sparse": p["is_sparse"],
+        }
+        self._ir_cache = d
+        return d
+
+    def param_count(self) -> dict:
+        """Return parameter-count estimates: total / active / per-layer breakdown."""
+        return self._params
+
+    def _repr_html_(self) -> str:
+        """Jupyter calls this; returned HTML string is rendered inline."""
+        return self._html(standalone=False)
+
+    def to_html(self, standalone: bool = True) -> str:
+        """Return the diagram as an HTML string.
+
+        Parameters
+        ----------
+        standalone : bool
+            If True (default), wraps the diagram in a full HTML document.
+            If False, returns a fragment usable for embedding (Jupyter mode).
+        """
+        return self._html(standalone=standalone)
+
+    def save(self, path: str) -> str:
+        """Save the diagram to disk.
+
+        - ``.html`` — interactive standalone document
+        - ``.json`` — the underlying IR (no rendering)
+        """
+        ext = os.path.splitext(path)[1].lower()
+        if ext == ".html":
+            with open(path, "w", encoding="utf-8") as f:
+                f.write(self.to_html(standalone=True))
+        elif ext == ".json":
+            with open(path, "w", encoding="utf-8") as f:
+                json.dump(self.to_ir(), f, indent=2)
+        else:
+            raise ValueError(
+                f"Unsupported extension {ext!r}. Use .html or .json."
+            )
+        return path
+
+    def _html(self, standalone: bool) -> str:
+        if standalone not in self._html_cache:
+            if standalone:
+                self._html_cache[standalone] = render_document(self.to_ir(), self._mount_id)
+            else:
+                self._html_cache[standalone] = render_fragment(self.to_ir(), self._mount_id)
+        return self._html_cache[standalone]
+
+    def __repr__(self) -> str:
+        return (
+            f"<Diagram {self.ir.name!r} · {self.ir.num_layers} layers · "
+            f"~{humanize(self._params['total'])} params"
+            + (f" ({humanize(self._params['active'])} active)" if self._params['is_sparse'] else "")
+            + ">"
+        )

model_unfolder/html_renderer.py

@@ -0,0 +1,5 @@
+"""Compatibility wrapper for the HTML/SVG renderer backend."""
+
+from .renderers.html import render_document, render_fragment
+
+__all__ = ["render_document", "render_fragment"]