Fix AttentionHead resizing bug (#63)

TransformerLensOrg · Aug 1, 2023 · 8605b16 · 8605b16
1 parent a4f6a72
commit 8605b16
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 20 deletions.
diff --git a/react/src/attention/AttentionHeads.tsx b/react/src/attention/AttentionHeads.tsx
@@ -143,17 +143,7 @@ export function AttentionHeads({
       <Row>
         <Col xs={12}>
           <h3 style={{ marginBottom: 10 }}>{headNames[focused]} Zoomed</h3>
-
-          <div
-            style={{
-              position: "relative",
-              // Set the maximum width such that a head with just a few tokens
-              // doesn't have crazy large boxes per token. Note this is the
-              // width of the full chart (including axis labels) so it also
-              // needs a sensible lowest maximum.
-              maxWidth: `${Math.max(Math.round(tokens.length * 2.4), 20)}em`
-            }}
-          >
+          <div>
             <h2
               style={{
                 position: "absolute",
@@ -174,6 +164,7 @@ export function AttentionHeads({
               minValue={minValue}
               negativeColor={negativeColor}
               positiveColor={positiveColor}
+              zoomed={true}
               tokens={tokens}
             />
           </div>

diff --git 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 (
     <Col>
-      <Row style={{ aspectRatio: showAxisLabels ? undefined : "1/1" }}>
-        <Chart
-          type="matrix"
-          options={options}
-          data={data}
-          width={1000}
-          height={1000}
-          updateMode="none"
-        />
+      <Row>
+        <div
+          style={{
+            // Chart.js charts resizing is weird.
+            // Responsive chart elements (which all are by default) require the
+            // parent element to have position: 'relative' and no sibling elements.
+            // There were previously issues that only occured at particular display
+            // sizes and zoom levels. See:
+            // https://github.com/alan-cooney/CircuitsVis/pull/63
+            // https://www.chartjs.org/docs/latest/configuration/responsive.html#important-note
+            // https://stackoverflow.com/a/48770978/7086623
+            position: "relative",
+            // Set the maximum width of zoomed heads such that a head with just a
+            // few tokens doesn't have crazy large boxes per token and the chart
+            // doesn't overflow the screen. Other heads fill their width.
+            maxWidth: zoomed
+              ? `min(100%, ${Math.round(tokens.length * 8)}em)`
+              : "initial",
+            width: zoomed ? "initial" : "100%",
+            aspectRatio: "1/1"
+          }}
+        >
+          <Chart
+            type="matrix"
+            options={options}
+            data={data}
+            width={1000}
+            height={1000}
+            updateMode="none"
+          />
+        </div>
       </Row>
     </Col>
   );
@@ -255,6 +278,11 @@ export interface AttentionPatternProps {
    */
   showAxisLabels?: boolean;
 
+  /**
+   * Is this a zoomed in view?
+   */
+  zoomed?: boolean;
+
   /**
    * List of tokens
    *