refactor(circuit): enhance node mapping and tracing result handling

Updated the map_idx_to_nodes method to include a target_layer argument for improved node mapping. Modified the extract_QK_tracing_result function to pass the layer parameter, ensuring accurate tracing results. Additionally, added a stringify_nodes method to the QKTracingResults class for better handling of node representations. Removed debug print statements in create_circuit function to clean up the code.

Author: Hzfinfdu

server/routers/circuits.py

@@ -139,8 +139,6 @@ def concretize_graph_data(graph_data: dict[str, Any]):
 def create_circuit(sae_set_name: str, request: GenerateCircuitRequest):
     """Generate and save a circuit graph for a given prompt and SAE set."""
 
-    print(request.model_dump())
-
     sae_set = client.get_sae_set(name=sae_set_name)
     assert sae_set is not None, f"SAE set {sae_set_name} not found"
     sae_names = sae_set.sae_names

src/lm_saes/circuit/utils/attn_scores_attribution.py

@@ -69,7 +69,7 @@ def drop_bias_terms(self):
         ), "We might not want to drop bias terms for the second time."
         return self
 
-    def map_idx_to_nodes(self, indices: torch.Tensor, input_ids: torch.Tensor) -> list[Node]:
+    def map_idx_to_nodes(self, indices: torch.Tensor, target_layer: int, input_ids: torch.Tensor) -> list[Node]:
         """Map component indices to node names.
 
         Args:
@@ -147,7 +147,7 @@ def map_idx_to_nodes(self, indices: torch.Tensor, input_ids: torch.Tensor) -> li
                 )
             else:
                 res = Node.bias_node(
-                    int(layer),
+                    target_layer.item(),
                     self.pos,
                     bias_name=self.bias_names[idx - decoder_bias_idx],
                     is_from_qk_tracing=True,
@@ -378,6 +378,7 @@ def get_single_side_QK_components(
 def extract_QK_tracing_result(
     q_side: ResStreamComponents,
     k_side: ResStreamComponents,
+    layer: int,
     input_ids: torch.Tensor,
     topk: int = 10,
 ) -> QKTracingResults:
@@ -402,8 +403,8 @@ def extract_QK_tracing_result(
     # pair-wise top contributors
     pair_wise_contributors = list(
         zip(
-            q_side.map_idx_to_nodes(q_features, input_ids),
-            k_side.map_idx_to_nodes(k_features, input_ids),
+            q_side.map_idx_to_nodes(q_features, layer, input_ids),
+            k_side.map_idx_to_nodes(k_features, layer, input_ids),
             topk_pairwise_attr_entries.values.cpu().tolist(),
         )
     )
@@ -413,13 +414,13 @@ def extract_QK_tracing_result(
     )
     top_q_marginal_contributors = list(
         zip(
-            q_side.map_idx_to_nodes(top_q_marginal_contributors.indices, input_ids),
+            q_side.map_idx_to_nodes(top_q_marginal_contributors.indices, layer, input_ids),
             top_q_marginal_contributors.values.cpu().tolist(),
         )
     )
     top_k_marginal_contributors = list(
         zip(
-            k_side.map_idx_to_nodes(top_k_marginal_contributors.indices, input_ids),
+            k_side.map_idx_to_nodes(top_k_marginal_contributors.indices, layer, input_ids),
             top_k_marginal_contributors.values.cpu().tolist(),
         )
     )
@@ -501,6 +502,6 @@ def compute_attn_scores_attribution(
         q_side=False,
     ).drop_bias_terms()
 
-    result = extract_QK_tracing_result(q_side, k_side, input_ids, topk=topk)
+    result = extract_QK_tracing_result(q_side, k_side, layer, input_ids, topk=topk)
     score_attribution_cache[cache_key] = result
     return result

src/lm_saes/circuit/utils/create_graph_files.py

@@ -163,13 +163,18 @@ def append_qk_tracing_results(graph: Graph, used_nodes: List[Node], clt_names, l
     for node in used_nodes:
         if node.qk_tracing_results is not None:
             from_qk_tracing_nodes.update(node.qk_tracing_results.get_nodes())
+
     nodes_to_add = from_qk_tracing_nodes - existing_nodes
     for node in nodes_to_add:
         if node.feature_type == "lorsa":
             node.sae_name = lorsa_names[node.layer // 2]
         elif node.feature_type == "cross layer transcoder":
             node.sae_name = clt_names[node.layer // 2]
         used_nodes.append(node)
+
+    for node in used_nodes:
+        if node.qk_tracing_results is not None:
+            node.qk_tracing_results.stringify_nodes()
     return used_nodes
 
 

src/lm_saes/circuit/utils/graph_file_utils.py

@@ -94,7 +94,7 @@ def bias_node(cls, layer, pos, bias_name, influence=None, is_from_qk_tracing=Fal
             node_id=f"{layer}_{bias_name}_{pos}",
             layer=layer,
             ctx_idx=pos,
-            feature_type=bias_name,
+            feature_type="bias",
             influence=influence,
             is_from_qk_tracing=is_from_qk_tracing,
         )
@@ -154,9 +154,11 @@ class QKTracingResults:
     top_k_marginal_contributors: List of tuples containing (k_node, score) for top K marginal contributors.
     """
 
-    pair_wise_contributors: list[tuple[Node, Node, float]]
-    top_q_marginal_contributors: list[tuple[Node, float]]
-    top_k_marginal_contributors: list[tuple[Node, float]]
+    NodeType = Node | str
+
+    pair_wise_contributors: list[tuple[NodeType, NodeType, float]]
+    top_q_marginal_contributors: list[tuple[NodeType, float]]
+    top_k_marginal_contributors: list[tuple[NodeType, float]]
 
     def get_nodes(self) -> set[Node]:
         all_relevant_nodes = set()
@@ -168,3 +170,30 @@ def get_nodes(self) -> set[Node]:
         for k_node, _ in self.top_k_marginal_contributors:
             all_relevant_nodes.add(k_node)
         return all_relevant_nodes
+
+    def stringify_nodes(self):
+        assert all(
+            isinstance(q_node, Node) and isinstance(k_node, Node) for q_node, k_node, _ in self.pair_wise_contributors
+        ) or all(
+            isinstance(q_node, str) and isinstance(k_node, str) for q_node, k_node, _ in self.pair_wise_contributors
+        )
+        assert all(isinstance(q_node, Node) for q_node, _ in self.top_q_marginal_contributors) or all(
+            isinstance(q_node, str) for q_node, _ in self.top_q_marginal_contributors
+        )
+        assert all(isinstance(k_node, Node) for k_node, _ in self.top_k_marginal_contributors) or all(
+            isinstance(k_node, str) for k_node, _ in self.top_k_marginal_contributors
+        )
+
+        if len(self.pair_wise_contributors) > 0 and isinstance(self.pair_wise_contributors[0][0], str):
+            return self
+
+        self.pair_wise_contributors = [
+            (q_node.node_id, k_node.node_id, score) for q_node, k_node, score in self.pair_wise_contributors
+        ]
+        self.top_q_marginal_contributors = [
+            (q_node.node_id, score) for q_node, score in self.top_q_marginal_contributors
+        ]
+        self.top_k_marginal_contributors = [
+            (k_node.node_id, score) for k_node, score in self.top_k_marginal_contributors
+        ]
+        return self