From bb06368602cfcca656746525de16a603e2359cb3 Mon Sep 17 00:00:00 2001
From: Ryan Mullins <ryanmullins@google.com>
Date: Thu, 28 Apr 2022 09:12:16 -0700
Subject: [PATCH] Adds GeneratedURL type and displays content to the
 GeneratedText and GeneratedImages modules.

PiperOrigin-RevId: 445168285
---
 lit_nlp/api/types.py                          |  8 +++++++-
 lit_nlp/client/lib/types.ts                   |  3 ++-
 .../client/modules/generated_image_module.ts  | 19 ++++++++++---------
 3 files changed, 19 insertions(+), 11 deletions(-)

diff --git a/lit_nlp/api/types.py b/lit_nlp/api/types.py
index d7713038..6a12f94f 100644
--- a/lit_nlp/api/types.py
+++ b/lit_nlp/api/types.py
@@ -195,7 +195,13 @@ class TopTokens(LitType):
 @attr.s(auto_attribs=True, frozen=True, kw_only=True)
 class URL(TextSegment):
   """TextSegment that should be interpreted as a URL."""
-  align: Text = None  # name of field in the model output
+  pass
+
+
+@attr.s(auto_attribs=True, frozen=True, kw_only=True)
+class GeneratedURL(TextSegment):
+  """A URL that was generated as part of a model prediction."""
+  align: Optional[Text] = None      # name of a field in the model output
 
 
 @attr.s(auto_attribs=True, frozen=True, kw_only=True)
diff --git a/lit_nlp/client/lib/types.ts b/lit_nlp/client/lib/types.ts
index 625c1e27..30041407 100644
--- a/lit_nlp/client/lib/types.ts
+++ b/lit_nlp/client/lib/types.ts
@@ -33,7 +33,8 @@ export type LitName = 'type'|'LitType'|'String'|'TextSegment'|'GeneratedText'|
     'AttentionHeads'|'SparseMultilabel'|'FieldMatcher'|'MultiFieldMatcher'|
     'Gradients'|'Boolean'|'TokenSalience'|'ImageBytes'|'SparseMultilabelPreds'|
     'ImageGradients'|'ImageSalience'|'SequenceSalience'|'ReferenceScores'|
-    'FeatureSalience'|'TopTokens'|'CurveDataPoints'|'InfluentialExamples';
+    'FeatureSalience'|'TopTokens'|'CurveDataPoints'|'InfluentialExamples'|
+    'GeneratedURL';
 
 export const listFieldTypes: LitName[] =
     ['Tokens', 'SequenceTags', 'SpanLabels', 'EdgeLabels', 'SparseMultilabel'];
diff --git a/lit_nlp/client/modules/generated_image_module.ts b/lit_nlp/client/modules/generated_image_module.ts
index 6211f118..cdb0ab4b 100644
--- a/lit_nlp/client/modules/generated_image_module.ts
+++ b/lit_nlp/client/modules/generated_image_module.ts
@@ -36,7 +36,7 @@ export class GeneratedImageModule extends LitModule {
   static override template = (model = '', selectionServiceIndex = 0) => {
     return html`
       <generated-image-module model=${model}
-        selectionServiceIndex=${selectionServiceIndex}>
+                              selectionServiceIndex=${selectionServiceIndex}>
       </generated-image-module>`;
   };
 
@@ -77,7 +77,8 @@ export class GeneratedImageModule extends LitModule {
     const dataset = this.appState.currentDataset;
     const promise = this.apiService.getPreds(
         [input], this.model, dataset,
-        [...GeneratedImageModule.supportedTypes, 'URL'], 'Generating images');
+        [...GeneratedImageModule.supportedTypes, 'GeneratedURL'],
+        'Generating images');
     const results = await this.loadLatest('generatedImages', promise);
     if (results === null) return;
 
@@ -113,17 +114,17 @@ export class GeneratedImageModule extends LitModule {
         <div class="module-results-area">
           ${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``;
+              const {align} = keyType;
+              if (isLitSubtype(keyType, 'GeneratedURL') &&
+                  align != null && align in this.generatedImages) {
+                return this.renderLink(key, value);
+              }
+
+              return null;
             })}
         </div>
       </div>