diff --git a/react/src/attention/AttentionHeads.tsx b/react/src/attention/AttentionHeads.tsx index f9c816c..6c601d0 100644 --- a/react/src/attention/AttentionHeads.tsx +++ b/react/src/attention/AttentionHeads.tsx @@ -143,17 +143,7 @@ export function AttentionHeads({

{headNames[focused]} Zoomed

- -
+

diff --git a/react/src/attention/AttentionPattern.tsx b/react/src/attention/AttentionPattern.tsx index 823428d..0bdcfd7 100644 --- a/react/src/attention/AttentionPattern.tsx +++ b/react/src/attention/AttentionPattern.tsx @@ -61,6 +61,7 @@ export function AttentionPattern({ positiveColor, upperTriColor = DefaultUpperTriColor, showAxisLabels = true, + zoomed = false, tokens }: AttentionPatternProps) { // Tokens must be unique (for the categories), so we add an index prefix @@ -164,15 +165,37 @@ export function AttentionPattern({ return ( - - + +
+ +
); @@ -255,6 +278,11 @@ export interface AttentionPatternProps { */ showAxisLabels?: boolean; + /** + * Is this a zoomed in view? + */ + zoomed?: boolean; + /** * List of tokens *