Remove attention module from LIT.

PiperOrigin-RevId: 644131463
PAIR-code · Jun 18, 2024 · 63665c8 · 63665c8
1 parent 0656386
commit 63665c8
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 358 deletions.
diff --git a/lit_nlp/api/layout.py b/lit_nlp/api/layout.py
@@ -32,7 +32,6 @@ class LitModuleName(dtypes.EnumSerializableAsValues, enum.Enum):
   # keep-sorted start
   AnnotatedTextGoldModule = 'annotated-text-gold-module'
   AnnotatedTextModule = 'annotated-text-module'
-  AttentionModule = 'attention-module'
   ClassificationModule = 'classification-module'
   ConfusionMatrixModule = 'confusion-matrix-module'
   CurvesModule = 'curves-module'
@@ -198,7 +197,6 @@ def to_json(self) -> dtypes.JsonDict:
         'Explanations': [
             modules.SalienceMapModule,
             modules.LegacySequenceSalienceModule,
-            modules.AttentionModule,
             modules.FeatureAttributionModule,
         ],
         'Clustering': [modules.SalienceClusteringModule],
@@ -232,7 +230,6 @@ def to_json(self) -> dtypes.JsonDict:
             *MODEL_PREDS_MODULES,
             modules.SalienceMapModule,
             modules.LegacySequenceSalienceModule,
-            modules.AttentionModule,
             modules.FeatureAttributionModule,
         ],
         'Salience Clustering': [modules.SalienceClusteringModule],

diff --git a/lit_nlp/client/modules/attention_module.ts b/lit_nlp/client/modules/attention_module.ts