feat(ui/circuit): show attention pattern

OpenMOSS · Sep 12, 2024 · 67bfb0b · 67bfb0b
1 parent 48a03a6
commit 67bfb0b
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 0 deletions.
diff --git a/ui/src/components/model/circuit.tsx b/ui/src/components/model/circuit.tsx
@@ -69,6 +69,10 @@ const NodeInfo = ({ node }: { node: Node<NodeData> }) => {
           <div className="text-sm">{node.data.tracingNode.key}</div>
           <div className="text-sm font-bold">Score:</div>
           <div className="text-sm">{node.data.tracingNode.activation.toFixed(3)}</div>
+          <div className="text-sm font-bold">Pattern:</div>
+          <div className={cn("text-sm", getAccentClassname(node.data.tracingNode.pattern, 1, "text"))}>
+            {node.data.tracingNode.pattern.toFixed(3)}
+          </div>
         </div>
       </div>
     );
@@ -211,6 +215,8 @@ export const CircuitViewer = memo(
     const getNodeClassNames = useCallback((node: TracingNode) => {
       if (node.type === "feature") {
         return cn(getAccentClassname(node.activation, node.maxActivation, "border"));
+      } else if (node.type === "attn-score") {
+        return cn(getAccentClassname(node.pattern, 1, "border"));
       }
       return "";
     }, []);

diff --git a/ui/src/types/model.ts b/ui/src/types/model.ts
@@ -26,6 +26,7 @@ const AttnScoreNodeSchema = z.object({
   query: z.number(),
   key: z.number(),
   activation: z.number(),
+  pattern: z.number(),
 });
 
 export const TracingNodeSchema = z.discriminatedUnion("type", [