update token colors

labmlai · Oct 19, 2024 · 60c6f35 · 60c6f35
1 parent 1a0aa78
commit 60c6f35
Show file tree

Hide file tree

Showing 3 changed files with 78 additions and 60 deletions.
diff --git a/ui/scss/charts.scss b/ui/scss/charts.scss
@@ -1,4 +1,7 @@
 .attention-visualization {
+    --text-color: #000000;
+    --background-color: #ffffff;
+
     color: #000000;
     background-color: #ffffff;
     .menu {
@@ -33,6 +36,9 @@
     }
 
     &.dark-theme {
+        --text-color: #ffffff;
+        --background-color: #111111;
+
         color: #ffffff;
         background-color: #111111;
 
@@ -121,7 +127,8 @@
             background-clip: padding-box !important;
 
             &.new-line {
-                opacity: 0.6;
+                opacity: 0.4;
+                font-size: 0.8em;
             }
 
             .underline {
@@ -133,6 +140,14 @@
                     width: 100%;
                 }
             }
+
+            border-top: 2px solid transparent;
+        }
+
+        .token:hover {
+            box-shadow: -4px 0px 12px rgba(0, 0, 0, 0.2),
+                4px 0px 12px rgba(0, 0, 0, 0.2);
+            border-top: 2px solid var(--text-color);
         }
 
         margin-bottom: 10px;
@@ -212,14 +227,10 @@
         left: auto;
     }
 
-    .hover-container:hover {
-        border: 1px solid #878686;
-    }
-
     .legend {
-        display: flex;
-        flex-direction: column;
-        width: fit-content;
+        display: grid;
+        grid-template-columns: 1fr 1fr;
+        gap: 10px;
 
         .row {
             display: flex;
@@ -230,26 +241,27 @@
             cursor: pointer;
             width: 100%;
             overflow-y: auto;
+        }
 
-            .legend-item {
-                display: flex;
-                align-items: center;
-                gap: 10px;
-            }
+        .legend-item {
+            display: flex;
+            align-items: center;
+            gap: 10px;
+            cursor: pointer;
+        }
 
-            .color {
-                width: 8px;
-                height: 8px;
-                border-radius: 50%;
-            }
+        .color {
+            width: 8px;
+            height: 8px;
+            border-radius: 50%;
+        }
 
-            .token-value {
-                font-family: monospace;
-                white-space: pre;
-                display: inline-block;
-                margin-left: 5px;
-                text-align: center;
-            }
+        .token-value {
+            font-family: monospace;
+            white-space: pre;
+            display: inline-block;
+            margin-left: 5px;
+            text-align: left;
         }
     }
 
@@ -260,7 +272,6 @@
         align-items: flex-start;
         justify-content: space-between;
         width: 100%;
-        max-width: 600px;
     }
 
     .selected-token {

diff --git a/ui/src/colors.ts b/ui/src/colors.ts
@@ -76,7 +76,7 @@ export class PlotColors {
             this.colorSchemes[chart as ChartType][
                 index % this.colorSchemes[chart as ChartType].length
             ].toLowerCase()
-        value = value * 0.6
+        value = value * 0.8
 
         return d3.interpolateRgb(col1, col2)(value)
     }

diff --git a/ui/src/token_loss.ts b/ui/src/token_loss.ts
@@ -170,14 +170,29 @@ export class StringTokenLoss {
         this.selectedTokenData = tokenData
         this.selectedTokenElem.textContent = this.selectedToken
         this.tokenInfoElem.textContent = tokenData.info
+
+        let value_strings = []
         for (let i = 0; i < tokenData.values.length; ++i) {
-            const valueName = tokenData.values[i].name
-            const valueElement = this.valueElements[valueName]
-            if (valueElement) {
-                valueElement.textContent = this.isScientific
+            value_strings.push(
+                this.isScientific
                     ? tokenData.values[i].value.toExponential()
                     : tokenData.values[i].value.toString()
-            }
+            )
+        }
+
+        const maxLength = Math.max(
+            ...value_strings.map((str) => str.split(".")[0].length)
+        )
+        value_strings = value_strings.map((str) => {
+            const [intPart, decPart] = str.split(".")
+            return (
+                intPart.padEnd(maxLength, " ") + (decPart ? "." + decPart : "")
+            )
+        })
+
+        for (let i = 0; i < value_strings.length; ++i) {
+            this.valueElements[tokenData.values[i].name].textContent =
+                value_strings[i]
         }
     }
 
@@ -239,41 +254,33 @@ export class StringTokenLoss {
                                 ++i
                             ) {
                                 let colorElem
-                                let row = $("div.row", ($) => {
-                                    $("span.legend-item", ($) => {
-                                        colorElem = $("div.color")
-                                        $(
-                                            "span",
-                                            this.tokenData[0].values[i].name
-                                        )
-                                    })
+                                let row = $("span.legend-item", ($) => {
+                                    colorElem = $("div.color")
+                                    $("span", this.tokenData[0].values[i].name)
+                                })
 
-                                    colorElem.style.setProperty(
-                                        "background",
+                                colorElem.style.setProperty(
+                                    "background",
+                                    this.colors.getInterpolatedColor(
+                                        1.0,
+                                        ChartType.TokenLoss,
+                                        i
+                                    )
+                                )
+                                colorElem.style.setProperty(
+                                    "border",
+                                    "2px solid " +
                                         this.colors.getInterpolatedColor(
                                             1.0,
                                             ChartType.TokenLoss,
                                             i
                                         )
-                                    )
-                                    colorElem.style.setProperty(
-                                        "border",
-                                        "2px solid " +
-                                            this.colors.getInterpolatedColor(
-                                                1.0,
-                                                ChartType.TokenLoss,
-                                                i
-                                            )
-                                    )
+                                )
 
-                                    let tokenValueElem = $(
-                                        "span.token-value",
-                                        ""
-                                    )
-                                    this.valueElements[
-                                        this.tokenData[0].values[i].name
-                                    ] = tokenValueElem as HTMLElement
-                                })
+                                let tokenValueElem = $("span.token-value", "")
+                                this.valueElements[
+                                    this.tokenData[0].values[i].name
+                                ] = tokenValueElem as HTMLElement
 
                                 row.addEventListener("click", () => {
                                     if (