Integrates a color legend into the sequence salience.

PiperOrigin-RevId: 468531709

lit_nlp/client/elements/color_legend.css

@@ -4,6 +4,7 @@
   align-items: center;
   height: 30px;
   overflow-x: auto;
+  overflow-y: hidden;
 }
 
 .legend-line {
@@ -15,7 +16,9 @@
 .legend-box {
   width: 13px;
   height: 13px;
-  margin: 5px;
+  margin: 2px;
+  border: 1px solid gray;
+  border-radius: 1px;
 }
 
 .legend-label {
@@ -56,4 +59,3 @@
   vertical-align: middle;
   margin: 0;
 }
-

lit_nlp/client/elements/color_legend.ts

@@ -66,7 +66,7 @@ export class ColorLegend extends ReactiveElement {
   private labelMarginRight: number = 2;
 
   private boxWidth: number = 13;
-  private boxMargin: number = 5;
+  private boxMargin: number = 2;
 
   private selectedColorLabelWidth: number = 46;
   private iconWidth: number = 16;
@@ -171,11 +171,10 @@ export class ColorLegend extends ReactiveElement {
     const domain = this.scale.domain();
     const hideLegend =
         domain.length === 1 && domain[0].toString().toLowerCase() === 'all';
-    const style = styleMap({'width': `${this.legendWidth}px`});
 
     // clang-format off
     return html`
-        <div class="legend-container" style=${style}>
+        <div class="legend-container">
           <mwc-icon class="palette-icon icon-outlined">palette</mwc-icon>
           <div class="color-label" title=${this.selectedColorName}
             name="color-name">
@@ -198,11 +197,9 @@ export class ColorLegend extends ReactiveElement {
     const maxValue = numDomain ? Math.max(...numDomain) : 0;
     const domain = linearSpace(minValue, maxValue, this.numBlocks || 5);
 
-    const style = styleMap({'width': `${this.legendWidth}px`});
-
     // clang-format off
     return html`
-        <div class="legend-container" style=${style}>
+        <div class="legend-container">
           <mwc-icon class="palette-icon icon-outlined">palette</mwc-icon>
           <div class="color-label" title=${this.selectedColorName}
             name="color-name">

lit_nlp/client/elements/slider.css

@@ -2,6 +2,7 @@
 input[type='range'].slider {
   -webkit-appearance: none;
   max-width: calc(100% - 4px);
+  min-width: 50px;
   background: var(--lit-neutral-400);
   border-radius: 4px;
   background-image: linear-gradient(var(--lit-cyea-400), var(--lit-cyea-400));

lit_nlp/client/modules/embeddings_module.css

@@ -17,9 +17,7 @@
   overflow: hidden;
 }
 
-.module-footer {
-  justify-content: space-between;
-}
+.module-footer > * { min-width: 0; }
 
 .selected-nearest-button:disabled {
   pointer-events: auto;

lit_nlp/client/modules/sequence_salience_module.css

@@ -9,6 +9,10 @@ table.field-table th {
   text-align: start;
 }
 
+lit-slider {
+  max-width: 72px;
+}
+
 .subfield-title {
   padding-top: 6px; /* align title with first row of content */
 }
@@ -158,27 +162,10 @@ table.field-table th {
   margin: 0 8px;
 }
 
-.cmap-blocks-holder {
-  display: flex;
-  flex-direction: row;
-  align-items: center;
-  margin: 0 4px;
-}
-
-.cmap-block {
-  width: 28px;
-  height: 20px;
-  border: 1px solid gray;
-  border-radius: 3px;
-  text-align: center;
-  vertical-align: middle;
-  margin: 0 2px;
-}
-
-.cmap-block-wide {
-  width: 32px;
-}
-
 .gamma-value {
   min-width: 32px;
 }
+
+.module-footer > * { min-width: 0; }
+
+.controls-group > * { min-width: 0; }

lit_nlp/client/modules/sequence_salience_module.ts

@@ -12,6 +12,7 @@ import {styleMap} from 'lit/directives/style-map';
 import {computed, observable} from 'mobx';
 
 import {LitModule} from '../core/lit_module';
+import {LegendType} from '../elements/color_legend';
 import {canonicalizeGenerationResults, GeneratedTextResult, GENERATION_TYPES, getAllOutputTexts, getAllReferenceTexts} from '../lib/generated_text_utils';
 import {Salience} from '../lib/lit_types';
 import {styles as sharedStyles} from '../lib/shared_styles.css';
@@ -361,29 +362,22 @@ export class SequenceSalienceModule extends LitModule {
     // clang-format on
   }
 
-  // TODO(b/198684817): move this colormap impl to a shared element?
-  renderColorBlocks() {
+  renderColorLegend() {
     const cmap = this.cmap;
     const isSigned = (cmap instanceof SignedSalienceCmap);
-    const blockValues = isSigned ? [-1.0, -0.7, -0.3, 0.0, 0.3, 0.7, 1.0] :
-                                   [0.0, 0.2, 0.4, 0.6, 0.8, 1.0];
+    const scale = (val: number) => cmap.bgCmap(val);
+    scale.domain = () => cmap.colorScale.domain();
+    const labelName = "Token Salience";
 
-    const blockClass = classMap({
-      'cmap-block': true,
-      'cmap-block-wide': isSigned,
-    });
-
-    const blocks = blockValues.map(val => {
-      const blockStyle = styleMap(
-          {'color': cmap.textCmap(val), 'background-color': cmap.bgCmap(val)});
-      // clang-format off
-      return html`
-        <div class=${blockClass} style=${blockStyle}>
-          ${val.toFixed(1)}
-        </div>`;
-      // clang-format on
-    });
-    return html`<div class='cmap-blocks-holder'>${blocks}</div>`;
+    // clang-format off
+    return html`
+        <color-legend legendType=${LegendType.SEQUENTIAL}
+          selectedColorName=${labelName}
+          ?alignRight=${true}
+          .scale=${scale}
+          numBlocks=${isSigned ? 7 : 5}>
+        </color-legend>`;
+    // clang-format on
   }
 
   renderFooterControls() {
@@ -413,8 +407,7 @@ export class SequenceSalienceModule extends LitModule {
         </div>
       </div>
       <div class="controls-group">
-        <label>Colormap:</label>
-        ${this.renderColorBlocks()}
+        ${this.renderColorLegend()}
         <label for="gamma-slider">Gamma:</label>
         <lit-slider min="0.25" max="6" step="0.25" val="${this.cmapGamma}"
                     .onInput=${onChangeGamma}></lit-slider>

lit_nlp/client/services/color_service.ts

@@ -33,15 +33,17 @@ export abstract class SalienceCmap {
    * An RGB interpolated color scale for one of the continuous LAB ramps from
    * VizColor, which have been linearized.
    */
-  protected colorScale: d3.ScaleSequential<string>;
+  protected myColorScale: d3.ScaleSequential<string>;
+
+  get colorScale() { return this.myColorScale; }
 
   // Exponent for computing luminance values from salience scores.
   // A higher value gives higher contrast for small (close to 0) salience
   // scores.
   // See https://en.wikipedia.org/wiki/Gamma_correction
   constructor(protected gamma: number = 1.0,
               protected domain: [number, number] = [0, 1]) {
-    this.colorScale = d3.scaleSequential(CONTINUOUS_UNSIGNED_LAB).domain(domain);
+    this.myColorScale = d3.scaleSequential(CONTINUOUS_UNSIGNED_LAB).domain(domain);
   }
 
   /**
@@ -54,7 +56,7 @@ export abstract class SalienceCmap {
 
   /** Clamps the value of d to the color scale's domain */
   clamp(d: number): number {
-    const [min, max] = this.colorScale.domain();
+    const [min, max] = this.myColorScale.domain();
     return Math.max(min, Math.min(max, d));
   }
 
@@ -75,20 +77,20 @@ export abstract class SalienceCmap {
 /** Color map for unsigned (positive) salience maps. */
 export class UnsignedSalienceCmap extends SalienceCmap {
   bgCmap(d: number): string {
-    return this.colorScale(this.lightness(d));
+    return this.myColorScale(this.lightness(d));
   }
 }
 
 /** Color map for signed salience maps. */
 export class SignedSalienceCmap extends SalienceCmap {
   constructor(gamma: number = 1.0, domain: [number, number] = [-1, 1]) {
     super(gamma, domain);
-    this.colorScale = d3.scaleSequential(CONTINUOUS_SIGNED_LAB).domain(domain);
+    this.myColorScale = d3.scaleSequential(CONTINUOUS_SIGNED_LAB).domain(domain);
   }
 
   bgCmap(d: number): string {
     const direction = d < 0 ? -1 : 1;
-    return this.colorScale(this.lightness(d) * direction);
+    return this.myColorScale(this.lightness(d) * direction);
   }
 }