From d69224b9357423c235bec1af80ef9cdb25f96f63 Mon Sep 17 00:00:00 2001
From: Cibi Arjun <arjuncibi@google.com>
Date: Fri, 20 Oct 2023 15:16:15 -0700
Subject: [PATCH] Internal Change

PiperOrigin-RevId: 575330002
---
 lit_nlp/api/dtypes.py                         |  7 +++++++
 lit_nlp/api/types.py                          |  7 +++++++
 lit_nlp/client/lib/dtypes.ts                  |  5 +++++
 lit_nlp/client/lib/lit_types.ts               | 10 +++++++++-
 lit_nlp/client/modules/salience_map_module.ts |  4 ++--
 5 files changed, 30 insertions(+), 3 deletions(-)

diff --git a/lit_nlp/api/dtypes.py b/lit_nlp/api/dtypes.py
index dbb6df82..dc0829ed 100644
--- a/lit_nlp/api/dtypes.py
+++ b/lit_nlp/api/dtypes.py
@@ -106,6 +106,13 @@ class TokenSalience(DataTuple):
   salience: Sequence[float]  # parallel to tokens
 
 
+@attr.s(auto_attribs=True, frozen=True, slots=True)
+class FrameSalience(DataTuple):
+  """Dataclass for a salience map over frames in a video."""
+
+  salience: dict[str, float]
+
+
 @attr.s(auto_attribs=True, frozen=True, slots=True)
 class FeatureSalience(DataTuple):
   """Dataclass for a salience map over categorical and/or scalar features."""
diff --git a/lit_nlp/api/types.py b/lit_nlp/api/types.py
index f3ea60c5..6ad5ec60 100644
--- a/lit_nlp/api/types.py
+++ b/lit_nlp/api/types.py
@@ -859,6 +859,13 @@ class ImageSalience(Salience):
   signed: bool = False  # If the returned values are signed.
 
 
+@attr.s(auto_attribs=True, frozen=True, kw_only=True)
+class FrameSalience(Salience):
+  """Metadata about a returned frame salience map."""
+
+  default: dtypes.FrameSalience = None
+
+
 @attr.s(auto_attribs=True, frozen=True, kw_only=True)
 class SequenceSalience(Salience):
   """Metadata about a returned sequence salience map."""
diff --git a/lit_nlp/client/lib/dtypes.ts b/lit_nlp/client/lib/dtypes.ts
index 9ab6acd9..c0cb6109 100644
--- a/lit_nlp/client/lib/dtypes.ts
+++ b/lit_nlp/client/lib/dtypes.ts
@@ -55,6 +55,11 @@ export interface TokenSalience extends DataTuple {
   salience: number[];  // Parallel to tokens.
 }
 
+/** Dataclass for a salience map over frames in a video. */
+export interface FrameSalience extends DataTuple {
+  salience: {[key: string]: number};
+}
+
 /** Dataclass for a salience map over categorical and/or scalar features. */
 export interface FeatureSalience extends DataTuple {
   salience: {[key: string]: number};
diff --git a/lit_nlp/client/lib/lit_types.ts b/lit_nlp/client/lib/lit_types.ts
index c334301e..dd442975 100644
--- a/lit_nlp/client/lib/lit_types.ts
+++ b/lit_nlp/client/lib/lit_types.ts
@@ -20,7 +20,7 @@
 // their Python counterparts.
 // tslint:disable:no-new-decorators class-name enforce-name-casing
 
-import {AnnotationCluster, EdgeLabel, FeatureSalience as FeatureSalienceDType, ScoredTextCandidates, SequenceSalienceMap, SpanLabel, TokenSalience as TokenSalienceDType} from './dtypes';
+import {AnnotationCluster, EdgeLabel, FeatureSalience as FeatureSalienceDType, ScoredTextCandidates, SequenceSalienceMap, SpanLabel, TokenSalience as TokenSalienceDType, FrameSalience as FrameSalienceDType} from './dtypes';
 
 /**
  * A dictionary of registered LitType names mapped to their constructor.
@@ -534,6 +534,14 @@ export class TokenSalience extends Salience {
   override default: TokenSalienceDType|undefined = undefined;
 }
 
+/**
+ * Metadata about a returned frame salience map.
+ */
+@registered
+export class FrameSalience extends Salience {
+  override default: FrameSalienceDType |undefined = undefined;
+}
+
 /**
  * Metadata about a returned feature salience map.
  */
diff --git a/lit_nlp/client/modules/salience_map_module.ts b/lit_nlp/client/modules/salience_map_module.ts
index 68992c98..73d4c325 100644
--- a/lit_nlp/client/modules/salience_map_module.ts
+++ b/lit_nlp/client/modules/salience_map_module.ts
@@ -35,7 +35,7 @@ import {LitModule} from '../core/lit_module';
 import {LegendType} from '../elements/color_legend';
 import {InterpreterClick} from '../elements/interpreter_controls';
 import {TokenWithWeight} from '../elements/token_chips';
-import {FeatureSalience, FieldMatcher, ImageGradients, ImageSalience, LitTypeTypesList, LitTypeWithParent, MulticlassPreds, RegressionScore, Salience, SalienceTargetInfo, TokenGradients, TokenSalience} from '../lib/lit_types';
+import {FeatureSalience, FieldMatcher, ImageGradients, ImageSalience, FrameSalience, LitTypeTypesList, LitTypeWithParent, MulticlassPreds, RegressionScore, Salience, SalienceTargetInfo, TokenGradients, TokenSalience} from '../lib/lit_types';
 import {styles as sharedStyles} from '../lib/shared_styles.css';
 import {CallConfig, IndexedInput, ModelInfoMap, Preds, SCROLL_SYNC_CSS_CLASS, Spec} from '../lib/types';
 import {argmax, cloneSpec, findSpecKeys, makeModifiedInput} from '../lib/utils';
@@ -69,7 +69,7 @@ type SalienceResult = TokenSalienceResult | ImageSalienceResult |
 
 // Notably, not SequenceSalience as that is handled by a different module.
 const SUPPORTED_SALIENCE_TYPES =
-    [TokenSalience, FeatureSalience, ImageSalience];
+    [TokenSalience, FeatureSalience, ImageSalience, FrameSalience];
 
 const TARGET_SELECTOR_SUPPORTED_TYPES: LitTypeTypesList =
     [MulticlassPreds, RegressionScore];