Promotes metrics info to top-level property of LitMetada.

PiperOrigin-RevId: 530374619

Author: RyanMullins

lit_nlp/app.py

@@ -20,7 +20,7 @@
 import random
 import threading
 import time
-from typing import Optional, Mapping, Sequence, Union, Callable, Iterable
+from typing import Callable, Iterable, Optional, Mapping, Sequence, TypedDict, Union
 
 from absl import logging
 
@@ -55,6 +55,38 @@
 ModelLoadersMap = dict[str, ModelLoader]
 
 
+# LINT.IfChange
+class ComponentInfo(TypedDict):
+  configSpec: types.Spec  # pylint: disable=invalid-name  # Named for JSON struct
+  metaSpec: types.Spec    # pylint: disable=invalid-name  # Named for JSON struct
+  description: str
+# LINT.ThenChange(./client/lib/types.ts)
+
+
+def _get_component_info(
+    obj: lit_components.Interpreter,
+) -> ComponentInfo:
+  """Returns the ComponentInfo for an Interpreter, Generator, Metric, etc."""
+  return ComponentInfo(
+      configSpec=obj.config_spec(),
+      metaSpec=obj.meta_spec(),
+      description=obj.description(),
+  )
+
+
+def _get_compatible_names(
+    candidates: Mapping[str, lit_components.Interpreter],
+    model: lit_model.Model,
+    dataset: lit_dataset.Dataset,
+) -> Sequence[str]:
+  """Returns the names of the candidates compatible with the model/dataset."""
+  return [
+      name
+      for name, candidate in candidates.items()
+      if candidate.is_compatible(model=model, dataset=dataset)
+  ]
+
+
 class LitApp(object):
   """LIT WSGI application."""
 
@@ -80,22 +112,29 @@ def _build_metadata(self):
 
       compat_gens: set[str] = set()
       compat_interps: set[str] = set()
+      compat_metrics: set[str] = set()
 
       for d in info['datasets']:
         dataset: lit_dataset.Dataset = self._datasets[d]
-        compat_gens.update([
-            name for name, gen in self._generators.items()
-            if gen.is_compatible(model=model, dataset=dataset)
-        ])
-        compat_interps.update([
-            name for name, interp in self._interpreters.items()
-            if interp.is_compatible(model=model, dataset=dataset)
-        ])
-
-      info['generators'] = [name for name in self._generators.keys()
-                            if name in compat_gens]
-      info['interpreters'] = [name for name in self._interpreters.keys()
-                              if name in compat_interps]
+        compat_gens.update(
+            _get_compatible_names(self._generators, model, dataset)
+        )
+        compat_interps.update(
+            _get_compatible_names(self._interpreters, model, dataset)
+        )
+        compat_metrics.update(
+            _get_compatible_names(self._metrics, model, dataset)
+        )
+
+      info['generators'] = [
+          name for name in self._generators.keys() if name in compat_gens
+      ]
+      info['interpreters'] = [
+          name for name in self._interpreters.keys() if name in compat_interps
+      ]
+      info['metrics'] = [
+          name for name in self._metrics.keys() if name in compat_metrics
+      ]
       model_info[name] = info
 
     dataset_info = {}
@@ -106,21 +145,19 @@ def _build_metadata(self):
           'size': len(ds),
       }
 
-    generator_info = {}
-    for name, gen in self._generators.items():
-      generator_info[name] = {
-          'configSpec': gen.config_spec(),
-          'metaSpec': gen.meta_spec(),
-          'description': gen.description()
-      }
+    generator_info: Mapping[str, ComponentInfo] = {
+        name: _get_component_info(gen) for name, gen in self._generators.items()
+    }
 
-    interpreter_info = {}
-    for name, interpreter in self._interpreters.items():
-      interpreter_info[name] = {
-          'configSpec': interpreter.config_spec(),
-          'metaSpec': interpreter.meta_spec(),
-          'description': interpreter.description()
-      }
+    interpreter_info: Mapping[str, ComponentInfo] = {
+        name: _get_component_info(interp)
+        for name, interp in self._interpreters.items()
+    }
+
+    metrics_info: Mapping[str, ComponentInfo] = {
+        name: _get_component_info(metric)
+        for name, metric in self._metrics.items()
+    }
 
     init_specs = {
         'datasets': {n: s for n, (_, s) in self._dataset_loaders.items()},
@@ -133,6 +170,7 @@ def _build_metadata(self):
         'datasets': dataset_info,
         'generators': generator_info,
         'interpreters': interpreter_info,
+        'metrics': metrics_info,
         'layouts': self._layouts,
         # Global configuration
         'demoMode': self._demo_mode,
@@ -569,6 +607,7 @@ def __init__(
       datasets: Mapping[str, lit_dataset.Dataset],
       generators: Optional[Mapping[str, lit_components.Generator]] = None,
       interpreters: Optional[Mapping[str, lit_components.Interpreter]] = None,
+      metrics: Optional[Mapping[str, lit_components.Metrics]] = None,
       annotators: Optional[list[lit_components.Annotator]] = None,
       layouts: Optional[layout.LitComponentLayouts] = None,
       dataset_loaders: Optional[DatasetLoadersMap] = None,
@@ -657,6 +696,11 @@ def __init__(
     else:
       self._interpreters = core.default_interpreters(self._models)
 
+    if metrics is not None:
+      self._metrics = metrics
+    else:
+      self._metrics = core.default_metrics()
+
     # Component to sync state from TS -> Python. Used in notebooks.
     if sync_state:
       self.ui_state_tracker = ui_state.UIStateTracker()

lit_nlp/client/lib/testing_utils.ts

@@ -85,7 +85,8 @@ export const mockMetadata: LitMetadata = {
       'generators':
           ['word_replacer', 'scrambler', 'backtranslation', 'hotflip'],
       'interpreters':
-          ['grad_norm', 'grad_sum', 'lime', 'metrics', 'pca', 'umap']
+          ['grad_norm', 'grad_sum', 'lime', 'metrics', 'pca', 'umap'],
+      'metrics': []
     },
     'sst_1_micro': {
       'spec': {
@@ -118,7 +119,8 @@ export const mockMetadata: LitMetadata = {
       'generators':
           ['word_replacer', 'scrambler', 'backtranslation', 'hotflip'],
       'interpreters':
-          ['grad_norm', 'grad_sum', 'lime', 'metrics', 'pca', 'umap']
+          ['grad_norm', 'grad_sum', 'lime', 'metrics', 'pca', 'umap'],
+      'metrics': []
     }
   },
   'datasets': {
@@ -182,6 +184,7 @@ export const mockMetadata: LitMetadata = {
     'pca': emptySpec(),
     'umap': emptySpec(),
   },
+  'metrics': {},
   'initSpecs': {
     'datasets': {
       'sst_dev': {'split': createLitType(StringLitType)},
@@ -250,7 +253,8 @@ export const mockSerializedMetadata: SerializedLitMetadata = {
       'generators':
           ['word_replacer', 'scrambler', 'backtranslation', 'hotflip'],
       'interpreters':
-          ['grad_norm', 'grad_sum', 'lime', 'metrics', 'pca', 'umap']
+          ['grad_norm', 'grad_sum', 'lime', 'metrics', 'pca', 'umap'],
+      'metrics': []
     },
     'sst_1_micro': {
       'spec': {
@@ -296,7 +300,8 @@ export const mockSerializedMetadata: SerializedLitMetadata = {
       'generators':
           ['word_replacer', 'scrambler', 'backtranslation', 'hotflip'],
       'interpreters':
-          ['grad_norm', 'grad_sum', 'lime', 'metrics', 'pca', 'umap']
+          ['grad_norm', 'grad_sum', 'lime', 'metrics', 'pca', 'umap'],
+      'metrics': []
     }
   },
   'datasets': {
@@ -375,6 +380,7 @@ export const mockSerializedMetadata: SerializedLitMetadata = {
     'pca': emptySpec(),
     'umap': emptySpec(),
   },
+  'metrics': {},
   'initSpecs': {
     'datasets': {
       'sst_dev': {

lit_nlp/client/lib/types.ts

@@ -42,11 +42,13 @@ interface InitSpecMap {
   [name: string]: Spec|null;  // using null here because None ==> null in Python
 }
 
+// LINT.IfChange
 export interface ComponentInfo {
   configSpec: Spec;
   metaSpec: Spec;
   description?: string;
 }
+// LINT.ThenChange(../../app.py)
 
 export interface DatasetInfo {
   size: number;
@@ -76,6 +78,7 @@ export interface ModelInfo {
   datasets: string[];
   generators: string[];
   interpreters: string[];
+  metrics: string[];
   spec: ModelSpec;
   description?: string;
 }
@@ -89,6 +92,7 @@ export interface LitMetadata {
   datasets: DatasetInfoMap;
   generators: ComponentInfoMap;
   interpreters: ComponentInfoMap;
+  metrics: ComponentInfoMap;
   layouts: LitComponentLayouts;
   demoMode: boolean;
   defaultLayout: string;

lit_nlp/client/modules/metrics_module.ts

@@ -299,6 +299,7 @@ export class MetricsModule extends LitModule {
   /** Convert the metricsMap information into table data for display. */
   @computed
   get tableData(): TableHeaderAndData {
+    // TODO(b/254832560): Use this.appState.metadata.metrics here.
     const {metaSpec} = this.appState.metadata.interpreters['metrics'];
     if (metaSpec == null) return {'header': [], 'data': []};
 

lit_nlp/client/services/classification_service_test.ts

@@ -85,7 +85,8 @@ describe('classification service test', () => {
         spec,
         datasets: [],
         generators: [],
-        interpreters: []
+        interpreters: [],
+        metrics: []
       }}
     };
 
@@ -165,7 +166,8 @@ describe('classification service test', () => {
           spec,
           datasets: [],
           generators: [],
-          interpreters: []
+          interpreters: [],
+          metrics: []
         }}
       };
 

lit_nlp/components/core.py

@@ -51,13 +51,8 @@
   _UMAP_AVAILABLE = False
 # pytype: enable=import-error # pylint: enable=g-import-not-at-top
 
-ComponentGroup = lit_components.ComponentGroup
-Generator = lit_components.Generator
-Interpreter = lit_components.Interpreter
-Model = lit_model.Model
 
-
-def default_generators() -> dict[str, Generator]:
+def default_generators() -> dict[str, lit_components.Generator]:
   """Returns a dict of the default generators used in a LitApp."""
   return {
       'Ablation Flip': ablation_flip.AblationFlip(),
@@ -67,21 +62,23 @@ def default_generators() -> dict[str, Generator]:
   }
 
 
-def required_interpreters() -> dict[str, Interpreter]:
+def required_interpreters() -> dict[str, lit_components.Interpreter]:
   """Returns a dict of required interpreters.
 
   These are used by multiple core modules, and without them the frontend will
   likely throw errors.
   """
   # Ensure the prediction analysis interpreters are included.
-  prediction_analysis_interpreters: dict[str, Interpreter] = {
+  prediction_analysis_interpreters: dict[str, lit_components.Interpreter] = {
       'classification': classification_results.ClassificationInterpreter(),
       'regression': regression_results.RegressionInterpreter(),
   }
   return prediction_analysis_interpreters
 
 
-def default_interpreters(models: dict[str, Model]) -> dict[str, Interpreter]:
+def default_interpreters(
+    models: dict[str, lit_model.Model]
+) -> dict[str, lit_components.Interpreter]:
   """Returns a dict of the default interpreters (and metrics) used in a LitApp.
 
   Args:
@@ -91,7 +88,7 @@ def default_interpreters(models: dict[str, Model]) -> dict[str, Interpreter]:
   interpreters = required_interpreters()
 
   # Ensure the embedding-based interpreters are included.
-  embedding_interpreters: dict[str, Interpreter] = {
+  embedding_interpreters: dict[str, lit_components.Interpreter] = {
       'nearest neighbors': nearest_neighbors.NearestNeighbors(),
       # Embedding projectors expose a standard interface, but get special
       # handling so we can precompute the projections if requested.
@@ -103,21 +100,23 @@ def default_interpreters(models: dict[str, Model]) -> dict[str, Interpreter]:
         umap.UmapModel
     )
 
-  gradient_map_interpreters: dict[str, Interpreter] = {
+  gradient_map_interpreters: dict[str, lit_components.Interpreter] = {
       'Grad L2 Norm': gradient_maps.GradientNorm(),
       'Grad ⋅ Input': gradient_maps.GradientDotInput(),
       'Integrated Gradients': gradient_maps.IntegratedGradients(),
       'LIME': lime_explainer.LIME(),
   }
 
   # pyformat: disable
-  core_interpreters: dict[str, Interpreter] = {
+  core_interpreters: dict[str, lit_components.Interpreter] = {
       'Model-provided salience': model_salience.ModelSalience(models),
       'counterfactual explainer': lemon_explainer.LEMON(),
       'tcav': tcav.TCAV(),
       'curves': curves.CurvesInterpreter(),
       'thresholder': thresholder.Thresholder(),
-      'metrics': default_metrics(),
+      # TODO(b/254832560): Remove this "metrics" record from the core
+      # interpreters once the front-end Metrics module has been updated.
+      'metrics': lit_components.ComponentGroup(default_metrics()),
       'pdp': pdp.PdpInterpreter(),
       'Salience Clustering': salience_clustering.SalienceClustering(
           dict(gradient_map_interpreters)),
@@ -133,13 +132,15 @@ def default_interpreters(models: dict[str, Model]) -> dict[str, Interpreter]:
   return interpreters
 
 
-def default_metrics() -> ComponentGroup:
-  return ComponentGroup({
+# TODO(b/254833485): Update typing to be a dict[str, lit_components.Metrics]
+# once the Wrapper classes in metrics.py inherit from lit_components.Metrics.
+def default_metrics() -> dict[str, lit_components.Interpreter]:
+  return {
       'regression': metrics.RegressionMetrics(),
       'multiclass': metrics.MulticlassMetrics(),
       'multilabel': metrics.MultilabelMetrics(),
       'paired': metrics.MulticlassPairedMetrics(),
       'bleu': metrics.CorpusBLEU(),
       'rouge': metrics.RougeL(),
       'exactmatch': metrics.ExactMatchMetrics(),
-  })
+  }

lit_nlp/components/metrics.py

@@ -152,6 +152,8 @@ def run_with_metadata(self,
     return ret
 
 
+# TODO(b/254833485): Convert to inherit from lit_components.Metrics so that
+# promotion of Metrics to a top-level class more direct.
 class ClassificationMetricsWrapper(lit_components.Interpreter):
   """Wrapper for classification metrics interpreters.