Create bridges for every module in neo (#987)

* Create bridge for every module in neo

* Use path to access subcomponents in get_component instead of Bridge name attribute values

* Handle remote names with multiple parts

---------

Co-authored-by: Bryce Meyer <[email protected]>

Author: degenfabian

transformer_lens/model_bridge/generalized_components/base.py

@@ -182,7 +182,15 @@ def _getattr_helper(self, name: str) -> Any:
         original_component = self._modules.get("_original_component", None)
         if original_component is not None:
             try:
-                return getattr(original_component, name)
+                name_split = name.split(".")
+
+                if len(name_split) > 1:
+                    current = getattr(original_component, name_split[0])
+                    for part in name_split[1:]:
+                        current = getattr(current, part)
+                    return current
+                else:
+                    return getattr(original_component, name)
             except AttributeError:
                 pass
 

transformer_lens/model_bridge/supported_architectures/neo.py

@@ -11,6 +11,7 @@
     AttentionBridge,
     BlockBridge,
     EmbeddingBridge,
+    LinearBridge,
     MLPBridge,
     NormalizationBridge,
     UnembeddingBridge,
@@ -66,9 +67,23 @@ def __init__(self, cfg: Any) -> None:
                 name="transformer.h",
                 submodules={
                     "ln1": NormalizationBridge(name="ln_1"),
-                    "attn": AttentionBridge(name="attn"),
+                    "attn": AttentionBridge(
+                        name="attn.attention",
+                        submodules={
+                            "W_Q": LinearBridge(name="q_proj"),
+                            "W_K": LinearBridge(name="k_proj"),
+                            "W_V": LinearBridge(name="v_proj"),
+                            "W_O": LinearBridge(name="out_proj"),
+                        },
+                    ),
                     "ln2": NormalizationBridge(name="ln_2"),
-                    "mlp": MLPBridge(name="mlp"),
+                    "mlp": MLPBridge(
+                        name="mlp",
+                        submodules={
+                            "W_in": LinearBridge(name="c_fc"),
+                            "W_out": LinearBridge(name="c_proj"),
+                        },
+                    ),
                 },
             ),
             "ln_final": NormalizationBridge(name="transformer.ln_f"),