Switch to pinning a datapoint for comparing datapoints.

- Adds pin button in datatable for pinning and unpinning. - Updates the UI around what was called the reference example, to be the pinned example, in the toolbar and in the duplicated module chrome. PiperOrigin-RevId: 436448987
PAIR-code · Mar 22, 2022 · 05bfc90 · 05bfc90
1 parent 7e80ba5
commit 05bfc90
Show file tree

Hide file tree

Showing 12 changed files with 154 additions and 65 deletions.
diff --git a/lit_nlp/client/core/app.ts b/lit_nlp/client/core/app.ts
@@ -77,23 +77,28 @@ export class LitApp {
       new Map<Constructor<LitService>, LitService|LitService[]>();
 
   /** Sync selection services */
-  syncSelectionServices() {
+  private syncSelectionServices() {
     const selectionServices = this.getServiceArray(SelectionService);
     // TODO(lit-dev): can we just copy the object instead, and skip this
     // logic?
     selectionServices[1].syncFrom(selectionServices[0]);
   }
 
   /** Simple DI service system */
-  getService<T extends LitService>(t: Constructor<T>): T {
+  getService<T extends LitService>(t: Constructor<T>, instance?: string): T {
     let service = this.services.get(t);
     /**
      * Modules that don't support example comparison will always get index
      * 0 of selectionService. This way we do not have to edit any module that
-     * does not explicitly support cloning
+     * does not explicitly support cloning. For modules that support comparison,
+     * if the `pinned` instance is specified then return the appropriate
+     * instance.
      */
     if (Array.isArray(service)) {
-      service = service[0];
+      if (instance != null && instance !== 'pinned') {
+        throw new Error(`Invalid service instance name: ${instance}`);
+      }
+      service = service[instance === 'pinned' ? 1 : 0];
     }
     if (service === undefined) {
       throw new Error(`Service is undefined: ${t.name}`);

diff --git a/lit_nlp/client/core/main_toolbar.css b/lit_nlp/client/core/main_toolbar.css
@@ -66,6 +66,7 @@
   display: flex;
   flex-direction: row;
   align-items: center;
+  margin-left: 4px;
 }
 
 #primary-selection-status {
@@ -170,3 +171,16 @@
   display: inline-block;
 }
 
+.pin-button {
+  width: 190px;
+}
+
+.pin-button-content {
+  display: flex;
+  width: 100%;
+}
+
+.pin-button-text {
+  flex-grow: 1;
+  text-align: center;
+}
diff --git a/lit_nlp/client/core/main_toolbar.ts b/lit_nlp/client/core/main_toolbar.ts
@@ -27,8 +27,7 @@ import {MobxLitElement} from '@adobe/lit-mobx';
 import {customElement} from 'lit/decorators';
 import { html} from 'lit';
 import {classMap} from 'lit/directives/class-map';
-import {styleMap} from 'lit/directives/style-map';
-import {computed, observable} from 'mobx';
+import {computed} from 'mobx';
 
 import {MenuItem} from '../elements/menu';
 import {styles as sharedStyles} from '../lib/shared_styles.css';
@@ -223,8 +222,6 @@ export class LitMainToolbar extends MobxLitElement {
   private readonly selectionService = app.getService(SelectionService);
   private readonly sliceService = app.getService(SliceService);
 
-  @observable private displayTooltip: boolean = false;
-
   /**
    * ID pairs, as [child, parent], from current selection or the whole dataset.
    */
@@ -358,7 +355,8 @@ export class LitMainToolbar extends MobxLitElement {
     const idPairs = this.selectedIdPairs;
     if (idPairs.length === 0) return null;
 
-    const referenceSelectionService = app.getServiceArray(SelectionService)[1];
+    const referenceSelectionService =
+        app.getService(SelectionService, 'pinned');
     const maybeGetIndex = (id: string|null) =>
         id != null ? this.appState.indicesById.get(id) : null;
     // ID pairs, as [child, parent]
@@ -450,47 +448,49 @@ export class LitMainToolbar extends MobxLitElement {
     const numSelected = this.selectionService.selectedIds.length;
     const numTotal = this.appState.currentInputData.length;
     const primaryId = this.selectionService.primarySelectedId;
+    const primaryIndex = primaryId == null ? -1 :
+        this.appState.indicesById.get(primaryId)!;
 
     const toggleExampleComparison = () => {
       this.appState.compareExamplesEnabled =
           !this.appState.compareExamplesEnabled;
     };
-    const compareDisabled =
-        !this.appState.compareExamplesEnabled && numSelected === 0;
-    const compareTextClass = classMap(
-        {'toggle-example-comparison': true, 'text-disabled': compareDisabled});
-    const tooltipStyle =
-        styleMap({visibility: this.displayTooltip ? 'visible' : 'hidden'});
-    const disableClicked = () => {
-      this.displayTooltip = true;
-    };
-    const disableMouseout = () => {
-      this.displayTooltip = false;
-    };
+
+    let referenceSelectedIndex = -1;
+    if (this.appState.compareExamplesEnabled) {
+      const referenceSelectionService =
+          app.getServiceArray(SelectionService)[1];
+      referenceSelectedIndex =
+          this.appState.indicesById.get(referenceSelectionService.primarySelectedId!)!;
+    }
+    const title = this.appState.compareExamplesEnabled ?
+        `Unpin datapoint ${referenceSelectedIndex}` :
+        primaryId == null ?
+            "Pin selected datapoint" : `Pin datapoint ${primaryIndex}`;
+    const pinDisabled = !this.appState.compareExamplesEnabled &&
+        primaryId == null;
+    const pinClasses = classMap({
+      'material-icon': true,
+      'span-outlined': !this.appState.compareExamplesEnabled
+    });
+    const buttonClasses = classMap({
+      'hairline-button': true,
+      'xl': true,
+      'pin-button': true,
+      'active': this.appState.compareExamplesEnabled
+    });
     // clang-format off
     return html`
     <div class='toolbar main-toolbar'>
       <div id='left-container'>
         <lit-main-menu></lit-main-menu>
-        <div class='compare-container'
-          @click=${compareDisabled ? null : toggleExampleComparison}
-          @mouseover=${compareDisabled ? disableClicked : null}
-          @mouseout=${compareDisabled ? disableMouseout : null}>
-          <div class=${compareTextClass}>
-            Compare datapoints
+        <button class="${buttonClasses}" title=${title}
+                ?disabled=${pinDisabled} @click=${toggleExampleComparison}>
+          <div class="pin-button-content">
+            <span class="${pinClasses}">push_pin</span>
+            <div class="pin-button-text">${title}</div>
           </div>
-          <mwc-switch id='compare-switch'
-            ?disabled='${compareDisabled}'
-            .checked=${this.appState.compareExamplesEnabled}
-          >
-          </mwc-switch>
-        </div>
-        <div class='compare-tooltip tooltip' style=${tooltipStyle}>
-          <mwc-icon class='tooltip-icon'>error_outline</mwc-icon>
-          <span class='tooltip-text'>
-            Select a datapoint to use this feature.
-          </span>
-        </div>
+        </button>
         ${this.renderPairControls()}
       </div>
       <div id='right-container'>

diff --git a/lit_nlp/client/core/widget.css b/lit_nlp/client/core/widget.css
@@ -42,17 +42,11 @@
   display: flex;
   flex-direction: row;
   align-items: center;
-  justify-content: space-between;
   padding: 0pt 4pt;
   border-bottom: 1px solid rgb(218, 220, 224);
   box-sizing: border-box;
 }
 
-.holder.highlight {
-  border: 1px solid #2f8c9b;
-  box-sizing: border-box;
-}
-
 .subtitle {
   font-size: 8pt;
   color: gray;
@@ -68,3 +62,8 @@
   align-items: center;
   justify-content: center;
 }
+
+.pin-spacer {
+  width: 16px;
+  height: 14px;
+}
diff --git a/lit_nlp/client/core/widget_group.ts b/lit_nlp/client/core/widget_group.ts
@@ -206,13 +206,13 @@ export class WidgetGroup extends LitElement {
 
     let subtitle = modelName ?? '';
     /**
-     * If defined, modules show "Main" for 0 and "Reference for 1,
+     * If defined, modules show "Selected" for 0 and "Pinned" for 1,
      * If undefined, modules do not show selectionService related info in their
      * titles (when compare examples mode is disabled)."
      */
     if (typeof selectionServiceIndex !== 'undefined' && moduleType.duplicateForExampleComparison) {
       subtitle = subtitle.concat(`${subtitle ? ' - ' : ''} ${
-        selectionServiceIndex ? 'Reference' : 'Main'}`);
+        selectionServiceIndex ? 'Pinned' : 'Selected'}`);
     }
     // Track scolling changes to the widget and request a rerender.
     const widgetScrollCallback = (event: CustomEvent<WidgetScroll>) => {
@@ -322,7 +322,6 @@ export class LitWidget extends MobxLitElement {
     });
     const holderClasses = classMap({
       holder: true,
-      highlight: this.highlight,
     });
     // Track content scrolling and pass the scrolling information back to the
     // widget group for sync'ing between duplicated widgets. This covers the
@@ -366,6 +365,10 @@ export class LitWidget extends MobxLitElement {
   renderHeader() {
     return html`
     <div class=header>
+      ${this.highlight ?
+        html`<mwc-icon class='material-icon pin-icon'>push_pin</mwc-icon>` :
+        html`<div class="pin-spacer"></div>`
+      }
       <span class="subtitle">${this.subtitle}</span>
     </div>
     `;

diff --git a/lit_nlp/client/elements/table.css b/lit_nlp/client/elements/table.css
@@ -149,10 +149,6 @@ tbody tr.focused {
   transition-property: background-color;
 }
 
-tbody tr.reference-selected {
-  outline: 1px solid #2f8c9b;
-}
-
 tbody td {
   vertical-align: top;
   min-width: 80px;

diff --git a/lit_nlp/client/elements/table.ts b/lit_nlp/client/elements/table.ts
@@ -784,7 +784,7 @@ export class DataTable extends ReactiveElement {
       'reference-selected': isReferenceSelection,
       'focused': isFocused
     });
-    const mouseDown = (e: MouseEvent) => {
+    const onClick = (e: MouseEvent) => {
       if (!this.selectionEnabled) return;
       this.handleRowClick(e, dataIndex, displayDataIndex);
     };
@@ -822,7 +822,7 @@ export class DataTable extends ReactiveElement {
     });
     // clang-format off
     return html`
-      <tr class="${rowClass}" @mousedown=${mouseDown} @mouseenter=${mouseEnter}
+      <tr class="${rowClass}" @click=${onClick} @mouseenter=${mouseEnter}
         @mouseleave=${mouseLeave}>
         ${data.rowData.map((d, i) =>
             html`<td style=${cellStyles}><div class=${cellClasses[i]}>${

diff --git a/lit_nlp/client/lib/shared_styles.css b/lit_nlp/client/lib/shared_styles.css
@@ -432,12 +432,27 @@ mwc-textfield {
   border: 1px solid var(--lit-neutral-200);
 }
 
-.hairline-button:active {
+.hairline-button:active,
+.hairline-button.active {
   background-color: var(--lit-mintonal-p-4);
   border: 1px solid var(--lit-cyea-500);
   color: var(--lit-cyea-500);
 }
 
+.pin-icon {
+  height: 14px;
+  width: 14px;
+  min-width: 14px;
+  --mdc-icon-size: 14px;
+  user-select: none;
+  color: var(--lit-cyea-500);
+  margin-right: 2px;
+}
+
+.icon-button.cyea {
+  color: var(--lit-cyea-500);
+}
+
 /**
  * For standalone MWC icons as buttons. We don't use mwc-icon-button because it
  * adds a large backdrop and extra whitespace.
@@ -470,6 +485,10 @@ mwc-icon.mdi-outlined {
   --mdc-icon-font: "Material Icons Outlined";
 }
 
+.span-outlined {
+  font-family: "Material Icons Outlined";
+}
+
 /**
  * Main container to use inside modules.
  * If this contains a .module-toolbar and a .module-results-area,

diff --git a/lit_nlp/client/modules/data_table_module.ts b/lit_nlp/client/modules/data_table_module.ts
@@ -21,6 +21,8 @@ import '../elements/checkbox';
 
 import {html} from 'lit';
 import {customElement, query} from 'lit/decorators';
+import {classMap} from 'lit/directives/class-map';
+import {styleMap} from 'lit/directives/style-map';
 import {computed, observable} from 'mobx';
 
 import {app} from '../core/app';
@@ -160,6 +162,11 @@ export class DataTableModule extends LitModule {
   // it gets run _four_ times every time a new datapoint is added.
   @computed
   get tableData(): TableData[] {
+    const pinnedId = this.appState.compareExamplesEnabled ?
+        app.getService(SelectionService, 'pinned').primarySelectedId : null;
+    const selectedId = this.selectionService.primarySelectedId;
+    const focusedId = this.focusService.focusData?.datapointId;
+
     // TODO(b/160170742): Make data table render immediately once the
     // non-prediction data is available, then fetch predictions asynchronously
     // and enable the additional columns when ready.
@@ -198,7 +205,51 @@ export class DataTableModule extends LitModule {
               .map(k => formatForDisplay(this.dataService.getVal(d.id, k),
                                          this.dataSpec[k]));
 
-      const ret: TableData = [index];
+      const pinClick = (event: Event) => {
+        if (pinnedId === d.id) {
+          this.appState.compareExamplesEnabled = false;
+        } else {
+          this.appState.compareExamplesEnabled = true;
+          app.getService(SelectionService, 'pinned').selectIds([d.id]);
+        }
+        event.stopPropagation();
+      };
+
+      const indexHolderDivStyle = styleMap({
+        'display': 'flex',
+        'flex-direction': 'row-reverse',
+        'justify-content': 'space-between',
+        'width': '100%'
+      });
+      const indexDivStyle = styleMap({
+        'text-align': 'right',
+      });
+      // Render the pin button next to the index if datapoint is pinned,
+      // selected, or hovered.
+      const renderPin = () => {
+        const iconClass = classMap({
+           'icon-button': true,
+           'cyea': true,
+           'mdi-outlined': pinnedId !== d.id,
+        });
+        if (pinnedId === d.id || focusedId === d.id || selectedId === d.id) {
+          return html`
+              <mwc-icon class="${iconClass}" @click=${pinClick}>
+                  push_pin
+              </mwc-icon>`;
+        }
+        return null;
+      };
+      const indexHtml = html`
+          <div style="${indexHolderDivStyle}">
+            <div style="${indexDivStyle}">${index}</div>
+            ${renderPin()}
+          </div>`;
+      const indexEntry = {
+        template: indexHtml,
+        value: index
+      };
+      const ret: TableData = [indexEntry];
       if (this.columnVisibility.get('id')) {
         ret.push(displayId);
       }
@@ -475,7 +526,7 @@ export class DataTableModule extends LitModule {
     let referenceSelectedIndex = -1;
     if (this.appState.compareExamplesEnabled) {
       const referenceSelectionService =
-          app.getServiceArray(SelectionService)[1];
+          app.getService(SelectionService, 'pinned');
       referenceSelectedIndex =
           indexOfId(referenceSelectionService.primarySelectedId);
     }