model-unfolder 0.2.4__tar.gz → 0.2.5__tar.gz

{model_unfolder-0.2.4 → model_unfolder-0.2.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: model-unfolder
-Version: 0.2.4
+Version: 0.2.5
 Summary: Unfold any HuggingFace transformer into an interactive architecture diagram, inline in Jupyter.
 Author: model-unfolder contributors
 License: Apache-2.0
@@ -111,8 +111,9 @@ Param estimates are close to published numbers — DeepSeek-V3 reports `~675B (~
 | DeepSeek | DeepSeek-V2, DeepSeek-V3, Kimi K2 |
 | Llama | Llama 3 / 3.1 / 3.2 / 3.3 |
 | Mistral | Mistral 7B, Mixtral 8x7B / 8x22B, Mistral Medium 3.5 |
-| Qwen | Qwen2 / 2.5, Qwen2-MoE, Qwen3, Qwen3-MoE, Qwen3.5 / 3.6 (hybrid linear+full attn) |
-| Gemma | Gemma 3, Gemma 4 (31B, E2B, E4B) |
+| Qwen | Qwen2 / 2.5, Qwen2-MoE, Qwen3, Qwen3-MoE, Qwen3.5 / 3.6 |
+| Gemma | Gemma 3 / 3n, Gemma 4 (31B, E2B, E4B) |
+| MiniMax | MiniMax-Text-01 |
 
 ### Diffusors
 

{model_unfolder-0.2.4 → model_unfolder-0.2.5}/README.md

@@ -87,8 +87,9 @@ Param estimates are close to published numbers — DeepSeek-V3 reports `~675B (~
 | DeepSeek | DeepSeek-V2, DeepSeek-V3, Kimi K2 |
 | Llama | Llama 3 / 3.1 / 3.2 / 3.3 |
 | Mistral | Mistral 7B, Mixtral 8x7B / 8x22B, Mistral Medium 3.5 |
-| Qwen | Qwen2 / 2.5, Qwen2-MoE, Qwen3, Qwen3-MoE, Qwen3.5 / 3.6 (hybrid linear+full attn) |
-| Gemma | Gemma 3, Gemma 4 (31B, E2B, E4B) |
+| Qwen | Qwen2 / 2.5, Qwen2-MoE, Qwen3, Qwen3-MoE, Qwen3.5 / 3.6 |
+| Gemma | Gemma 3 / 3n, Gemma 4 (31B, E2B, E4B) |
+| MiniMax | MiniMax-Text-01 |
 
 ### Diffusors
 

model_unfolder-0.2.5/model_unfolder/adapters/transformer/families/__init__.py

@@ -0,0 +1,14 @@
+"""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, fallback, llama, minimax, mistral, qwen
+from . import gemma  # gemma/ package — dispatches to gemma3/gemma4 internally
+
+# Order matters: more specific adapters first.
+# ``gemma`` must run before ``llama`` (llama previously caught "gemma" model_type).
+# ``mistral``, ``qwen``, ``minimax`` before ``llama``.
+# ``fallback`` always matches — must be last.
+ADAPTERS = [deepseek, gemma, minimax, mistral, qwen, llama, fallback]

model_unfolder-0.2.5/model_unfolder/adapters/transformer/families/gemma/__init__.py

@@ -0,0 +1,25 @@
+"""Gemma family adapter — dispatches to version-specific sub-adapters.
+
+Sub-adapters (in priority order):
+  gemma4  — Gemma 4 (31B, E2B, E4B, 26B-A4B)   model_type: gemma4 / gemma4_text
+  gemma3  — Gemma 3 / 3n (1B–27B, E2B, E4B)     model_type: gemma3 / gemma3n
+"""
+from __future__ import annotations
+
+from typing import Any
+
+from . import gemma3, gemma4
+
+_SUB_ADAPTERS = [gemma4, gemma3]
+
+
+def matches(cfg: Any) -> bool:
+    return any(a.matches(cfg) for a in _SUB_ADAPTERS)
+
+
+def parse(cfg: Any) -> Any:
+    for a in _SUB_ADAPTERS:
+        if a.matches(cfg):
+            return a.parse(cfg)
+    # Should not be reached since matches() already confirmed one hit
+    return gemma3.parse(cfg)

model_unfolder-0.2.5/model_unfolder/adapters/transformer/families/gemma/gemma3.py

@@ -0,0 +1,137 @@
+"""Adapter for Google's Gemma 3 and Gemma 3n families.
+
+Gemma 3 (1B / 4B / 12B / 27B) is a multimodal wrapper (model_type: "gemma3")
+nesting the language model under text_config, with alternating
+sliding-window and full-context attention layers controlled by
+``sliding_window_pattern`` (every Nth layer is full).
+
+Gemma 3n (E2B / E4B) is the nano variant (model_type: "gemma3n") with
+Per-Layer Embeddings — structurally similar but with PLE conditioning.
+"""
+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
+from ...special_parts.per_layer_embedding import (
+    per_layer_embedding_blocks,
+    per_layer_embedding_extras,
+)
+
+
+_TOP_TYPES  = {"gemma3", "gemma3n"}
+_TEXT_TYPES = {"gemma3_text", "gemma3n_text"}
+_ALL_TYPES  = _TOP_TYPES | _TEXT_TYPES
+_ARCH_HINTS = ("gemma3",)
+
+
+def matches(cfg: Any) -> bool:
+    model_type = (_g(cfg, "model_type") or "").lower()
+    if model_type in _ALL_TYPES:
+        return True
+    arches = _g(cfg, "architectures") or []
+    return any(any(h in a.lower() for h in _ARCH_HINTS) for a in arches)
+
+
+def parse(cfg: Any) -> ModelIR:
+    arch_name = architecture_name(cfg, "gemma3")
+    text_cfg  = _text_config(cfg)
+    model_type = (_g(text_cfg, "model_type") or _g(cfg, "model_type") or "").lower()
+
+    num_layers    = _g(text_cfg, "num_hidden_layers", 0)
+    hidden_size   = _g(text_cfg, "hidden_size", 0)
+    num_heads     = _g(text_cfg, "num_attention_heads", 0)
+    num_kv_heads  = _g(text_cfg, "num_key_value_heads", num_heads)
+    head_dim      = _g(text_cfg, "head_dim") or (hidden_size // num_heads if num_heads else None)
+    intermediate_size = _g(text_cfg, "intermediate_size", 0)
+    activation    = (_g(text_cfg, "hidden_activation") or _g(text_cfg, "hidden_act") or "gelu").lower()
+
+    if num_kv_heads == num_heads:
+        attn_kind = "mha"
+    elif num_kv_heads == 1:
+        attn_kind = "mqa"
+    else:
+        attn_kind = "gqa"
+
+    sliding_window  = _g(text_cfg, "sliding_window")
+    # Gemma 3: every sliding_window_pattern-th layer is full, rest are sliding
+    sliding_pattern = _g(text_cfg, "sliding_window_pattern")
+    layer_types     = _g(text_cfg, "layer_types") or []
+
+    # PLE (Gemma 3n)
+    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 = []
+    for i in range(num_layers):
+        if layer_types and i < len(layer_types):
+            lt = layer_types[i]
+            if "sliding" in lt:
+                mask, win = "sliding", sliding_window
+            else:
+                mask, win = "global", None
+        elif sliding_pattern and sliding_window:
+            is_full = (i % sliding_pattern) == (sliding_pattern - 1)
+            mask = "global" if is_full else "sliding"
+            win  = None if is_full else sliding_window
+        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,
+        )
+
+        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) or _g(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,
+    )
+
+    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,
+        extras=extras,
+    )
+
+
+def _text_config(cfg: Any) -> Any:
+    model_type = (_g(cfg, "model_type") or "").lower()
+    if model_type in _TOP_TYPES:
+        sub = _g(cfg, "text_config")
+        if sub is not None:
+            return sub
+    return cfg

{model_unfolder-0.2.4/model_unfolder/adapters/transformer/families → model_unfolder-0.2.5/model_unfolder/adapters/transformer/families/gemma}/gemma4.py

@@ -18,10 +18,10 @@ 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 (
+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,
 )

{model_unfolder-0.2.4 → model_unfolder-0.2.5}/model_unfolder/adapters/transformer/families/llama.py

@@ -8,7 +8,7 @@ from ..assembly import decoder_extras, decoder_layer
 from ..common import architecture_name, get_config_value as _g, model_name
 
 
-_FAMILIES = {"llama", "phi3", "gemma"}
+_FAMILIES = {"llama", "phi3"}
 
 
 def matches(cfg: Any) -> bool:

model_unfolder-0.2.5/model_unfolder/adapters/transformer/families/minimax.py

@@ -0,0 +1,99 @@
+"""Adapter for MiniMax models (MiniMax-Text-01 and successors).
+
+MiniMax-Text-01 is a 456B sparse-MoE hybrid:
+* ``attn_type_list`` — per-layer int array: 0 = lightning (linear) attention,
+  1 = full softmax attention.  Every 8th layer is full (7 linear + 1 full).
+* MoE FFN: ``num_local_experts`` routed experts, ``num_experts_per_tok`` active.
+* Flat config (no text_config wrapper).
+"""
+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
+
+
+_MODEL_TYPES = {"minimax_text_01", "minimax"}
+
+
+def matches(cfg: Any) -> bool:
+    model_type = (_g(cfg, "model_type") or "").lower()
+    if model_type in _MODEL_TYPES:
+        return True
+    arches = _g(cfg, "architectures") or []
+    return any("minimax" in a.lower() for a in arches)
+
+
+def parse(cfg: Any) -> ModelIR:
+    arch_name = architecture_name(cfg, "minimax")
+
+    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)
+    activation   = (_g(cfg, "hidden_act") or "silu").lower()
+
+    if num_kv_heads == num_heads:
+        attn_kind = "mha"
+    elif num_kv_heads == 1:
+        attn_kind = "mqa"
+    else:
+        attn_kind = "gqa"
+
+    # 0 = lightning/linear attention, 1 = full softmax attention
+    attn_type_list = _g(cfg, "attn_type_list") or []
+
+    # MoE
+    num_experts         = _g(cfg, "num_local_experts") or _g(cfg, "num_experts") or 0
+    num_experts_per_tok = _g(cfg, "num_experts_per_tok") or 0
+    intermediate_size   = _g(cfg, "intermediate_size", 0)
+    is_moe = bool(num_experts)
+
+    layers = []
+    for i in range(num_layers):
+        attn_flag = attn_type_list[i] if i < len(attn_type_list) else 1
+        is_linear = attn_flag == 0
+
+        attn = AttentionSpec(
+            kind="linear" if is_linear else attn_kind,
+            num_heads=num_heads,
+            num_kv_heads=num_kv_heads,
+            head_dim=head_dim,
+            mask="causal",
+        )
+
+        if is_moe:
+            ffn = FFNSpec(
+                kind="moe",
+                activation=activation,
+                intermediate_size=intermediate_size,
+                gated=True,
+                num_experts=num_experts,
+                num_experts_per_tok=num_experts_per_tok,
+                expert_intermediate_size=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),
+    )

{model_unfolder-0.2.4 → model_unfolder-0.2.5}/model_unfolder/renderers/html/sections.py

@@ -26,6 +26,16 @@ def _header(ir: dict, info: dict) -> str:
         badges.append(
             f'<span class="uf-badge" title="{_attr(title)}">{_html(badge["text"])}</span>'
         )
+
+    warnings = ir.get("warnings") or []
+    if warnings:
+        tooltip = " · ".join(warnings)
+        badges.append(
+            f'<span class="uf-badge uf-badge-warn" title="{_attr(tooltip)}">'
+            "⚠ partial config"
+            "</span>"
+        )
+
     return f"""
 <div class="uf-header">
   <div class="uf-name">{_html(ir.get("name", "model"))}</div>

{model_unfolder-0.2.4 → model_unfolder-0.2.5}/model_unfolder/renderers/html/styles.py

@@ -52,6 +52,12 @@ def _style(mount_id: str) -> str:
   font-weight:600;
   letter-spacing:0.02em;
 }}
+#{mount_id} .uf-badge-warn {{
+  background:#FEF3C7;
+  color:#92400E;
+  border:1px solid #FCD34D;
+  cursor:default;
+}}
 #{mount_id} .uf-stats {{
   display:grid;
   grid-template-columns:repeat(5,minmax(0,1fr));

{model_unfolder-0.2.4 → model_unfolder-0.2.5}/model_unfolder.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: model-unfolder
-Version: 0.2.4
+Version: 0.2.5
 Summary: Unfold any HuggingFace transformer into an interactive architecture diagram, inline in Jupyter.
 Author: model-unfolder contributors
 License: Apache-2.0
@@ -111,8 +111,9 @@ Param estimates are close to published numbers — DeepSeek-V3 reports `~675B (~
 | DeepSeek | DeepSeek-V2, DeepSeek-V3, Kimi K2 |
 | Llama | Llama 3 / 3.1 / 3.2 / 3.3 |
 | Mistral | Mistral 7B, Mixtral 8x7B / 8x22B, Mistral Medium 3.5 |
-| Qwen | Qwen2 / 2.5, Qwen2-MoE, Qwen3, Qwen3-MoE, Qwen3.5 / 3.6 (hybrid linear+full attn) |
-| Gemma | Gemma 3, Gemma 4 (31B, E2B, E4B) |
+| Qwen | Qwen2 / 2.5, Qwen2-MoE, Qwen3, Qwen3-MoE, Qwen3.5 / 3.6 |
+| Gemma | Gemma 3 / 3n, Gemma 4 (31B, E2B, E4B) |
+| MiniMax | MiniMax-Text-01 |
 
 ### Diffusors
 

{model_unfolder-0.2.4 → model_unfolder-0.2.5}/model_unfolder.egg-info/SOURCES.txt

@@ -23,10 +23,13 @@ model_unfolder/adapters/transformer/common.py
 model_unfolder/adapters/transformer/families/__init__.py
 model_unfolder/adapters/transformer/families/deepseek.py
 model_unfolder/adapters/transformer/families/fallback.py
-model_unfolder/adapters/transformer/families/gemma4.py
 model_unfolder/adapters/transformer/families/llama.py
+model_unfolder/adapters/transformer/families/minimax.py
 model_unfolder/adapters/transformer/families/mistral.py
 model_unfolder/adapters/transformer/families/qwen.py
+model_unfolder/adapters/transformer/families/gemma/__init__.py
+model_unfolder/adapters/transformer/families/gemma/gemma3.py
+model_unfolder/adapters/transformer/families/gemma/gemma4.py
 model_unfolder/adapters/transformer/special_parts/__init__.py
 model_unfolder/adapters/transformer/special_parts/per_layer_embedding.py
 model_unfolder/renderers/__init__.py

{model_unfolder-0.2.4 → model_unfolder-0.2.5}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "model-unfolder"
-version = "0.2.4"
+version = "0.2.5"
 description = "Unfold any HuggingFace transformer into an interactive architecture diagram, inline in Jupyter."
 readme = "README.md"
 requires-python = ">=3.9"

model_unfolder-0.2.4/model_unfolder/adapters/transformer/families/__init__.py

@@ -1,13 +0,0 @@
-"""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, fallback, gemma4, llama, mistral, qwen
-
-# Order matters: more specific adapters first.
-# ``gemma4`` must run before ``llama`` (which also matches ``gemma`` model_type).
-# ``mistral`` and ``qwen`` must run before ``llama``.
-# ``fallback`` always matches — must be last.
-ADAPTERS = [deepseek, gemma4, mistral, qwen, llama, fallback]