Adds clickable URLs to the Generated Images module.

PiperOrigin-RevId: 441987921
PAIR-code · Apr 15, 2022 · 8cf8119 · 8cf8119
1 parent c4e8606
commit 8cf8119
Show file tree

Hide file tree

Showing 2 changed files with 45 additions and 13 deletions.
diff --git a/lit_nlp/api/types.py b/lit_nlp/api/types.py
@@ -196,7 +196,7 @@ class TopTokens(LitType):
 @attr.s(auto_attribs=True, frozen=True, kw_only=True)
 class URL(TextSegment):
   """TextSegment that should be interpreted as a URL."""
-  pass
+  align: Text = None  # name of field in the model output
 
 
 @attr.s(auto_attribs=True, frozen=True, kw_only=True)

diff --git a/lit_nlp/client/modules/generated_image_module.ts b/lit_nlp/client/modules/generated_image_module.ts
@@ -23,7 +23,7 @@ import {observable} from 'mobx';
 import {LitModule} from '../core/lit_module';
 import {styles as sharedStyles} from '../lib/shared_styles.css';
 import {IndexedInput, LitName, ModelInfoMap, Spec} from '../lib/types';
-import {doesOutputSpecContain} from '../lib/utils';
+import {doesOutputSpecContain, isLitSubtype} from '../lib/utils';
 
 /**
  * A LIT module that renders generated text.
@@ -46,19 +46,28 @@ export class GeneratedImageModule extends LitModule {
   static override get styles() {
     const styles = css`
       .field-group {
-        padding: 4px;
+        padding: 8px;
       }
-    `;
+
+      .field-title {
+        color: var(--lit-bold-label-color);
+      }
+
+      .field-group a {
+        display: flex;
+        flex-direction: row;
+        align-items: center;
+        justify-content: left;
+      }`;
     return [sharedStyles, styles];
   }
 
   @observable private generatedImages: {[key: string]: string} = {};
 
   override firstUpdated() {
     this.reactImmediately(
-        () => this.selectionService.primarySelectedInputData, data => {
-          this.updateSelection(data);
-        });
+        () => this.selectionService.primarySelectedInputData,
+        data => {this.updateSelection(data);});
   }
 
   private async updateSelection(input: IndexedInput|null) {
@@ -67,32 +76,55 @@ export class GeneratedImageModule extends LitModule {
 
     const dataset = this.appState.currentDataset;
     const promise = this.apiService.getPreds(
-        [input], this.model, dataset, GeneratedImageModule.supportedTypes,
-        'Generating images');
+        [input], this.model, dataset,
+        [...GeneratedImageModule.supportedTypes, 'URL'], 'Generating images');
     const results = await this.loadLatest('generatedImages', promise);
     if (results === null) return;
 
     this.generatedImages = results[0];
   }
 
-  renderOutputGroup(name: string) {
+  renderImage(name: string, src: string) {
     // clang-format off
     return html`
       <div class='field-group'>
         <div class="field-title">${name}</div>
-        <img src=${this.generatedImages[name]}>
+        <img src=${src}>
       </div>
     `;
     // clang-format on
   }
 
+  renderLink(name: string, url: string) {
+    // clang-format off
+    return html`
+      <div class='field-group'>
+        <div class="field-title">${name}</div>
+        <a href=${url} target="_blank">${url}</a>
+      </div>`;
+    // clang-format on
+  }
+
   override render() {
+    const {output} = this.appState.getModelSpec(this.model);
     // clang-format off
     return html`
       <div class='module-container'>
         <div class="module-results-area">
-          ${Object.keys(this.generatedImages).map(
-              name => this.renderOutputGroup(name))}
+          ${Object.entries(this.generatedImages).map(([key, value]) => {
+              const keyType = output[key];
+
+              if (isLitSubtype(keyType, 'URL') && keyType.align != null &&
+                  keyType.align in this.generatedImages) {
+                return this.renderLink(key, value);
+              }
+
+              if (isLitSubtype(keyType, GeneratedImageModule.supportedTypes)) {
+                return this.renderImage(key, value);
+              }
+
+              return html``;
+            })}
         </div>
       </div>
     `;