Custom regex separator mode for token grouping in Sequence Salience.

PiperOrigin-RevId: 619324011
PAIR-code · Mar 26, 2024 · c8ee224 · c8ee224
1 parent b9941ed
commit c8ee224
Show file tree

Hide file tree

Showing 3 changed files with 120 additions and 10 deletions.
diff --git a/lit_nlp/client/elements/token_chips.css b/lit_nlp/client/elements/token_chips.css
@@ -5,6 +5,8 @@
 }
 
 .tokens-holder {
+  --tokens-line-height: 22px;
+
   display: flex;
   flex-direction: row;
   flex-wrap: wrap;
@@ -68,6 +70,14 @@
   height: 0;
 }
 
+/**
+ * If there are multiple row breaks, subsequent ones should create full blank
+ * lines.
+ */
+.row-break + .row-break {
+  height: calc(var(--tokens-line-height) + 3.5px);
+}
+
 .word-spacer {
   width: 1em;
 }
@@ -82,7 +92,7 @@
 .text-chips {
   display: block;
   font-size: 0; /* hack to get zero spacing between elements */
-  line-height: 22px;
+  line-height: var(--tokens-line-height);
 }
 
 .text-chips > * {
@@ -168,7 +178,8 @@
 
 /* vertical dense mode */
 .text-chips.vdense {
-  line-height: 16px;
+  --tokens-line-height: 16px;
+  /* line-height: 16px; */
 }
 .text-chips.vdense .salient-token {
   padding: 1.5px 0; /* avoid highlight area overlapping across lines */

diff --git a/lit_nlp/client/modules/lm_salience_module.css b/lit_nlp/client/modules/lm_salience_module.css
@@ -145,7 +145,9 @@ select:invalid {
 
 #change-target-icon {
   display: none; /* hide by default */
-  line-height: 36px; /* vertical alignment issue */
+  mwc-icon {
+    line-height: 36px; /* vertical alignment issue */
+  }
 }
 
 .vertical-separator {
@@ -176,14 +178,24 @@ select:invalid {
 /**
  * Container queries to hide some labels in narrow windows.
  * Helps in SxS mode when the module is replicated.
+ *
+ * Adjust the size of these so the labels are hidden in a graceful way as the
+ * module is resized:
+ * - Method label is on the right, hide this first as fewer things will jump.
+ * - Remaining things can hide at a narrower width
+ *
+ * Adjust sizes based on the actual width of toolbar elements so that labels
+ * will be hidden /before/ items wrap to the next row. For example, if adding
+ * an additional button that is 20px wide (including margins), add 20px to
+ * the widths so that the labels hide sooner.
  */
-@container (width < 720px) {
+@container (width < 745px) {
   #method-label {
     display: none;
   }
 }
 
-@container (width < 660px) {
+@container (width < 685px) {
   #change-target-button, #granularity-label, #colormap-slider-label {
     display: none;
   }
@@ -221,4 +233,24 @@ color-legend {
   width: 100%;
   height: 2px;
   animation: running-progress 2s cubic-bezier(0.4, 0, 0.2, 1) infinite;
-}
+}
+
+.regex-input-container {
+  display: flex;
+  flex-direction: row;
+  align-items: center;
+  gap: 2px;
+}
+
+.regex-input {
+  max-width: 100px;
+}
+
+.error-input {
+  border-color: var(--google-red-600);
+  border-style: solid;
+  border-radius: 3px;
+  border-width: 2px;
+  padding-top: 1px;
+  padding-bottom: 1px;
+}
diff --git a/lit_nlp/client/modules/lm_salience_module.ts b/lit_nlp/client/modules/lm_salience_module.ts
@@ -42,7 +42,7 @@ enum SegmentationMode {
   LINES = 'Lines',
   PARAGRAPHS = '¶',
   // TODO(b/324961811): add phrase or clause chunking?
-  // TODO(b/324961803): add custom regex?
+  CUSTOM = '⚙',
 }
 
 const LEGEND_INFO_TITLE_SIGNED =
@@ -202,6 +202,8 @@ const REQUEST_PENDING: unique symbol = Symbol('REQUEST_PENDING');
 
 const CMAP_DEFAULT_RANGE = 0.4;
 
+const DEFAULT_CUSTOM_SEGMENTATION_REGEX = '\\n+';
+
 /** LIT module for model output. */
 @customElement('lm-salience-module')
 export class LMSalienceModule extends SingleExampleSingleModelModule {
@@ -226,6 +228,9 @@ export class LMSalienceModule extends SingleExampleSingleModelModule {
 
   @observable
   private segmentationMode: SegmentationMode = SegmentationMode.WORDS;
+  @observable
+  private customSegmentationRegexString: string =
+      DEFAULT_CUSTOM_SEGMENTATION_REGEX;
   // TODO(b/324959547): get default from spec
   @observable private selectedSalienceMethod? = 'grad_l2';
   // Output range for colormap.
@@ -305,6 +310,17 @@ export class LMSalienceModule extends SingleExampleSingleModelModule {
     return makeModifiedInput(this.currentData, {'target': targetString});
   }
 
+  @computed
+  get customSegmentationRegex(): RegExp|undefined {
+    try {
+      return RegExp(this.customSegmentationRegexString, 'g');
+    } catch (error) {
+      console.warn(
+          'Invalid segmentation regex: ', this.customSegmentationRegexString);
+      return undefined;
+    }
+  }
+
   @computed
   get currentTokenGroups(): string[][] {
     if (this.segmentationMode === SegmentationMode.TOKENS) {
@@ -330,6 +346,14 @@ export class LMSalienceModule extends SingleExampleSingleModelModule {
     } else if (this.segmentationMode === SegmentationMode.PARAGRAPHS) {
       // Paragraph separator is two or more newlines.
       return groupTokensByRegexSeparator(this.currentTokens, /\n\n+/g);
+    } else if (this.segmentationMode === SegmentationMode.CUSTOM) {
+      if (this.customSegmentationRegex === undefined) {
+        // Just return tokens.
+        return this.currentTokens.map(t => [t]);
+      } else {
+        return groupTokensByRegexPrefix(
+            this.currentTokens, this.customSegmentationRegex);
+      }
     } else {
       throw new Error(
           `Unsupported segmentation mode ${this.segmentationMode}.`);
@@ -567,6 +591,10 @@ export class LMSalienceModule extends SingleExampleSingleModelModule {
       return {
         text: val,
         selected: this.segmentationMode === val,
+        tooltipText:
+            (val === SegmentationMode.PARAGRAPHS ? 'Paragraphs' :
+                 val === SegmentationMode.CUSTOM ? 'Custom Regex' :
+                                                   ''),
         onClick: () => {
           if (this.segmentationMode !== val) {
             this.resetTargetSpan();
@@ -584,6 +612,38 @@ export class LMSalienceModule extends SingleExampleSingleModelModule {
       this.underline = !this.underline;
     };
 
+    const updateSegmentationRegex = (e: Event) => {
+      const {value} = e.target as HTMLInputElement;
+      this.customSegmentationRegexString = value;
+      this.resetTargetSpan();
+    };
+
+    const regexEntryClasses = classMap({
+      'regex-input': true,
+      // Note: customSegmentationRegex is the RegExp object, it will be
+      // undefined if the customSegmentationRegexString does not define a valid
+      // regular expression.
+      'error-input': this.customSegmentationRegex === undefined
+    });
+
+    const resetSegmentationRegex = () => {
+      this.customSegmentationRegexString = DEFAULT_CUSTOM_SEGMENTATION_REGEX;
+    };
+
+    // prettier-ignore
+    const customRegexEntry = html`
+      <div class='regex-input-container'>
+        <input type='text' class=${regexEntryClasses} slot='tooltip-anchor'
+          title="Enter a regex to match segment prefix."
+          @input=${updateSegmentationRegex}
+          .value=${this.customSegmentationRegexString}>
+        <mwc-icon class='icon-button value-reset-icon' title='Reset regex'
+          @click=${resetSegmentationRegex}>
+          restart_alt
+        </mwc-icon>
+      </div>
+    `;
+
     // prettier-ignore
     return html`
       <div class="controls-group" style="gap: 8px;">
@@ -593,6 +653,9 @@ export class LMSalienceModule extends SingleExampleSingleModelModule {
             .options=${segmentationOptions}
             ?disabled=${this.currentTokens.length === 0}>
         </lit-fused-button-bar>
+        ${
+        this.segmentationMode === SegmentationMode.CUSTOM ? customRegexEntry :
+                                                            null}
       </div>
       <div class="controls-group" style="gap: 8px;">
         <lit-switch
@@ -724,13 +787,17 @@ export class LMSalienceModule extends SingleExampleSingleModelModule {
         </div>
       </div>
       <div class='controls-group'>
-        <lit-tooltip content=${targetSelectorHelp} tooltipPosition="left">
-          <button class='hairline-button' id='change-target-button'
+        <lit-tooltip content=${targetSelectorHelp} tooltipPosition="left"
+          id='change-target-button'>
+          <button class='hairline-button'
             slot='tooltip-anchor' @click=${clearSalienceTarget}
             ?disabled=${target == null}>
             <span>Select sequence </span><span class='material-icon'>arrow_drop_down</span>
           </button>
-          <mwc-icon class='icon-button' id='change-target-icon'
+        </lit-tooltip>
+        <lit-tooltip content=${targetSelectorHelp} tooltipPosition="left"
+          id='change-target-icon'>
+          <mwc-icon class='icon-button'
             slot='tooltip-anchor' @click=${clearSalienceTarget}>
             edit
           </mwc-icon>