model-unfolder 0.2.7__tar.gz → 0.2.8__tar.gz

{model_unfolder-0.2.7 → model_unfolder-0.2.8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: model-unfolder
-Version: 0.2.7
+Version: 0.2.8
 Summary: Unfold any HuggingFace transformer into an interactive architecture diagram, inline in Jupyter.
 Author: model-unfolder contributors
 License: Apache-2.0
@@ -95,7 +95,7 @@ No extra config in `model_unfolder` itself.
 ```python
 diagram = unfold(cfg)
 diagram.save("model.html")   # standalone interactive HTML
-diagram.save("model.json")   # IR (no rendering)
+diagram.save("model.json")   # expanded architecture JSON (no rendering)
 diagram.param_count()        # {"total": ..., "active": ..., "per_layer": [...]}
 diagram.to_ir()              # full IR dict
 ```

{model_unfolder-0.2.7 → model_unfolder-0.2.8}/README.md

@@ -71,7 +71,7 @@ No extra config in `model_unfolder` itself.
 ```python
 diagram = unfold(cfg)
 diagram.save("model.html")   # standalone interactive HTML
-diagram.save("model.json")   # IR (no rendering)
+diagram.save("model.json")   # expanded architecture JSON (no rendering)
 diagram.param_count()        # {"total": ..., "active": ..., "per_layer": [...]}
 diagram.to_ir()              # full IR dict
 ```

{model_unfolder-0.2.7 → model_unfolder-0.2.8}/model_unfolder/__init__.py

@@ -11,6 +11,7 @@ Outside Jupyter::
     diagram.save("kimi_k2.html")
 """
 from .diagram import Diagram
+from .evidence import inspect_model_code
 from .parser import config_to_ir
 from .ir import ModelIR, LayerSpec, AttentionSpec, FFNSpec, CrossLayerEdge
 from .params import estimate_params
@@ -27,11 +28,19 @@ __all__ = [
     "FFNSpec",
     "CrossLayerEdge",
     "config_to_ir",
+    "inspect_model_code",
     "estimate_params",
 ]
 
 
-def unfold(cfg_or_id, token=None) -> Diagram:
+def unfold(
+    cfg_or_id,
+    token=None,
+    *,
+    inspect_code: bool = False,
+    code_source: str = "local",
+    return_json: bool = False,
+):
     """Unfold a transformer into a renderable architecture diagram.
 
     Parameters
@@ -44,14 +53,31 @@ def unfold(cfg_or_id, token=None) -> Diagram:
         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.
+    inspect_code
+        If True, attach static source-code evidence to the IR. The code scanner
+        parses modeling files as text/AST and does not execute model code.
+    code_source
+        Source for code inspection: ``"local"`` (installed transformers),
+        ``"path"``, ``"hub"``, ``"auto"``, or a local file/directory path.
+    return_json
+        If True, return the expanded architecture JSON dict instead of the
+        renderable ``Diagram``.  The JSON uses stable structural fields for
+        dimensions, projections, layer groups, operation graphs, cache behavior,
+        and trace paths instead of renderer labels/descriptions.
 
     Returns
     -------
-    Diagram
-        Renders inline in Jupyter; otherwise call ``.save()`` or ``.to_html()``.
+    Diagram | dict
+        ``Diagram`` by default; ``dict`` when ``return_json=True``.
     """
-    ir = config_to_ir(cfg_or_id, token=token)
-    return Diagram(ir)
+    ir = config_to_ir(
+        cfg_or_id,
+        token=token,
+        inspect_code=inspect_code,
+        code_source=code_source,
+    )
+    diagram = Diagram(ir)
+    return diagram.to_json() if return_json else diagram
 
 
 # friendly alias

model_unfolder-0.2.8/model_unfolder/adapters/transformer/__init__.py

@@ -0,0 +1,8 @@
+"""Transformer-LLM adapter.
+
+There is exactly one parser (``parser.py``); see its module docstring for
+the principle (config-driven, no per-family code).
+"""
+from . import parser
+
+ADAPTERS = [parser]

{model_unfolder-0.2.7 → model_unfolder-0.2.8}/model_unfolder/adapters/transformer/blocks/attention.py

@@ -48,12 +48,18 @@ def _sdpa_child_blocks(attention: AttentionSpec, hidden_size: int) -> list[dict]
         {
             "id": "k_proj",
             "title": "Key projection",
-            "description": f"Linear; {hidden} -> {kv_out}  ({num_kv_heads} KV-heads x {d_k} dims)",
+            "description": (
+                f"Linear; {hidden} -> {kv_out}  ({num_kv_heads} KV-heads x {d_k} dims). "
+                "Cache ports show K/V write/read during generation: arrowhead for write, blunt tail for read."
+            ),
         },
         {
             "id": "v_proj",
             "title": "Value projection",
-            "description": f"Linear; {hidden} -> {kv_out}  ({num_kv_heads} KV-heads x {d_k} dims)",
+            "description": (
+                f"Linear; {hidden} -> {kv_out}  ({num_kv_heads} KV-heads x {d_k} dims). "
+                "Cache ports show K/V write/read during generation: arrowhead for write, blunt tail for read."
+            ),
         },
         {
             "id": "qkv_dot",
@@ -104,12 +110,18 @@ def _sdpa_detailed_child_blocks(
         {
             "id": "k_proj",
             "title": "Key projection",
-            "description": f"Linear; {hidden} -> {kv_out}  ({kv_label} x {d_k} dims)",
+            "description": (
+                f"Linear; {hidden} -> {kv_out}  ({kv_label} x {d_k} dims). "
+                "Cache ports show K/V write/read during generation: arrowhead for write, blunt tail for read."
+            ),
         },
         {
             "id": "v_proj",
             "title": "Value projection",
-            "description": f"Linear; {hidden} -> {kv_out}  ({kv_label} x {d_k} dims)",
+            "description": (
+                f"Linear; {hidden} -> {kv_out}  ({kv_label} x {d_k} dims). "
+                "Cache ports show K/V write/read during generation: arrowhead for write, blunt tail for read."
+            ),
         },
         {
             "id": "scaled_scores",
@@ -232,7 +244,10 @@ def _mla_child_blocks(attention: AttentionSpec, hidden_size: int) -> list[dict]:
             "id": "mla_cache",
             "label": "latent cache c_t",
             "title": "Stored latent cache",
-            "description": f"Compressed K/V latent stored in the cache instead of full K and V heads; rank {kv_rank}",
+            "description": (
+                f"Compressed K/V latent stored in the cache instead of full K and V heads; rank {kv_rank}. "
+                "Cache ports show write from compression and read back into K/V expansion."
+            ),
         },
         {
             "id": "mla_kv_up",
@@ -289,7 +304,7 @@ def _mla_child_blocks(attention: AttentionSpec, hidden_size: int) -> list[dict]:
             "title": "MLA K/V cache path",
             "description": (
                 f"Compresses hidden state into rank {kv_rank} latent cache, expands K/V content, "
-                "and combines K noPE with a RoPE key side-channel"
+                "and combines K noPE with a RoPE key side-channel. Cache ports mark the latent write/read point."
             ),
             "detail_view": "mla_kv_cache_path",
             "children": kv_children,

{model_unfolder-0.2.7 → model_unfolder-0.2.8}/model_unfolder/adapters/transformer/blocks/descriptions.py

@@ -67,13 +67,14 @@ def describe_attention(attention: AttentionSpec) -> str:
         )
         if attention.q_lora_rank:
             text += f"; Q LoRA {_fmt(attention.q_lora_rank)}"
+        text += "; cache ports mark latent write/read state"
         return text
     if attention.kind == "mqa":
-        return _with_attention_window(attention, f"Multi-query; {attention.num_heads} Q / 1 KV head")
+        return _with_attention_window(attention, f"Multi-query; {attention.num_heads} Q / 1 KV head; cache ports mark K/V write/read state")
     if attention.kind == "gqa":
         return _with_attention_window(attention, (
             f"Grouped-query; {attention.num_heads} Q / {attention.num_kv_heads} KV heads; "
-            f"head dim {_fmt(attention.head_dim)}"
+            f"head dim {_fmt(attention.head_dim)}; cache ports mark K/V write/read state"
         ))
     if attention.kind == "ssm":
         shared = "; weight-shared across positions" if attention.shared else ""
@@ -94,11 +95,12 @@ def describe_attention(attention: AttentionSpec) -> str:
     if attention.no_rope:
         extras.append("NoPE")
     suffix = f"; {', '.join(extras)}" if extras else ""
-    return _with_attention_window(attention, f"Multi-head; {attention.num_heads} heads; head dim {_fmt(attention.head_dim)}{suffix}")
+    cache_note = "; cache ports mark K/V write/read state"
+    return _with_attention_window(attention, f"Multi-head; {attention.num_heads} heads; head dim {_fmt(attention.head_dim)}{suffix}{cache_note}")
 
 
 def _attention_mask_prefix(attention: AttentionSpec) -> str:
-    return "SWA" if attention.mask == "sliding" else ""
+    return "SW" if attention.mask == "sliding" else ""
 
 
 def _attention_mask_title_prefix(attention: AttentionSpec) -> str: