class_docstr: Salience map clustering.
class_docstr_tok: ["Salience","map","clustering","."]
docstr: Run this component, given a model and input(s).
  
  Note that when `config['REUSE_CLUSTERING'] == True` we reuse the previously
  computed kmeans objects and vocabulary. Tokens that do not exist in the
  vocabulary will be ignored.
  
  Args:
    inputs: Inputs to cluster.
    model: Model that provides salience maps.
    dataset: Dataset to compute salience maps for.
    model_outputs: Precomputed model outputs.
    config: Config for clustering and salience computation
  
  Returns:
    Dict with 2 keys:
      `CLUSTER_ID_KEY`: Contains the cluster assignments. One cluster id per
        dataset example.
      `REPRESENTATION_KEY`: Contains the representations of all examples in
        the dataset that were used in the clustering.
      `TOP_TOKEN_KEY`: Top tokens are the tokens that have highest mean
        salience over all examples. They are sorted according to their
        salience.
  
  Raises:
    RuntimeError: Salience interpreter incompatible with model and dataset
    TypeError: config is not provided
    ValueError: config["Salience Mapper"] is not provided
docstr_tok: ["Run","this","component",",","given","a","model","and","input","(","s",")",".","Note","that","when","`","config","[","'REUSE_CLUSTERING","'","]","==","True","`","we","reuse","the","previously","computed","kmeans","objects","and","vocabulary",".","Tokens","that","do","not","exist","in","the","vocabulary","will","be","ignored",".","Args",":","inputs",":","Inputs","to","cluster",".","model",":","Model","that","provides","salience","maps",".","dataset",":","Dataset","to","compute","salience","maps","for",".","model_outputs",":","Precomputed","model","outputs",".","config",":","Config","for","clustering","and","salience","computation","Returns",":","Dict","with","2","keys",":","`","CLUSTER_ID_KEY","`",":","Contains","the","cluster","assignments",".","One","cluster","id","per","dataset","example",".","`","REPRESENTATION_KEY","`",":","Contains","the","representations","of","all","examples","in","the","dataset","that","were","used","in","the","clustering",".","`","TOP_TOKEN_KEY","`",":","Top","tokens","are","the","tokens","that","have","highest","mean","salience","over","all","examples",".","They","are","sorted","according","to","their","salience",".","Raises",":","RuntimeError",":","Salience","interpreter","incompatible","with","model","and","dataset","TypeError",":","config","is","not","provided","ValueError",":","config","[","``","Salience","Mapper","''","]","is","not","provided"]
code: def run(
      self,
      inputs: Sequence[JsonDict],
      model: lit_model.Model,
      dataset: lit_dataset.Dataset,
      model_outputs: Optional[list[JsonDict]] = None,
      config: Optional[JsonDict] = None,
  ) -> Optional[JsonDict]:
      """Run this component, given a model and input(s).
  
      Note that when `config['REUSE_CLUSTERING'] == True` we reuse the previously
      computed kmeans objects and vocabulary. Tokens that do not exist in the
      vocabulary will be ignored.
  
      Args:
        inputs: Inputs to cluster.
        model: Model that provides salience maps.
        dataset: Dataset to compute salience maps for.
        model_outputs: Precomputed model outputs.
        config: Config for clustering and salience computation
  
      Returns:
        Dict with 2 keys:
          `CLUSTER_ID_KEY`: Contains the cluster assignments. One cluster id per
            dataset example.
          `REPRESENTATION_KEY`: Contains the representations of all examples in
            the dataset that were used in the clustering.
          `TOP_TOKEN_KEY`: Top tokens are the tokens that have highest mean
            salience over all examples. They are sorted according to their
            salience.
  
      Raises:
        RuntimeError: Salience interpreter incompatible with model and dataset
        TypeError: config is not provided
        ValueError: config["Salience Mapper"] is not provided
      """
  
      if not config:
          raise TypeError("config must be provided")
  
      salience_key: Optional[str] = config.get(SALIENCE_MAPPER_KEY)
      if not salience_key:
          raise ValueError(
              f"config[{SALIENCE_MAPPER_KEY}] must be provided"
          )
  
      if not (
          salience_interpreter := self.salience_mappers.get(
              salience_key
          )
      ):
          raise ValueError(
              f"No interpreter registered for ${salience_key}."
          )
  
      if not salience_interpreter.is_compatible(
          model=model, dataset=dataset
      ):
          raise RuntimeError(
              f"The {salience_key} interpreter is incompatible with this model and"
              " dataset pair."
          )
  
      # If no specific inputs provided, use the entire dataset.
      inputs_to_use = inputs or dataset.examples
      token_saliencies = salience_interpreter.run(
          inputs_to_use, model, dataset, model_outputs, config
      )
  
      if not token_saliencies:
          return None
  
      salience_fields = list(token_saliencies[0].keys())
      reuse_clustering = self.kmeans and config.get(
          REUSE_CLUSTERING, False
      )
  
      if not reuse_clustering:
          try:
              self.vocab = model.get_embedding_table()[0]
              self.vocab_lookup = {
                  word_type: i for i, word_type in enumerate(self.vocab)
              }
          except NotImplementedError:
              self.vocab_lookup, self.vocab = self._build_vocab(
                  token_saliencies
              )
  
      top_k = int(
          config.get(
              TOP_K_TOKENS_KEY,
              self.config_spec()[TOP_K_TOKENS_KEY].default,
          )
      )
  
      representations = self._compute_fixed_length_representation(
          token_saliencies, self.vocab_lookup, top_k
      )
  
      cluster_ids = {}
      salience_field_to_representations = {}
      salience_field_to_top_tokens = {}
  
      for salience_field in salience_fields:
          weight_matrix = np.vstack(
              [
                  representation[salience_field]
                  for representation in representations
              ]
          )
          n_clusters = int(
              config.get(
                  N_CLUSTERS_KEY,
                  self.config_spec()[N_CLUSTERS_KEY].default,
              )
          )
          seed = int(
              config.get(SEED_KEY, self.config_spec()[SEED_KEY].default)
          )
          if not reuse_clustering:
              self.kmeans[salience_field] = cluster.KMeans(
                  n_clusters=n_clusters, random_state=seed
              )
              cluster_ids[salience_field] = (
                  self.kmeans[salience_field]
                  .fit_predict(weight_matrix)
                  .tolist()
              )
          else:
              cluster_ids[salience_field] = (
                  self.kmeans[salience_field]
                  .predict(weight_matrix)
                  .tolist()
              )
  
          salience_field_to_representations[
              salience_field
          ] = weight_matrix
          salience_field_to_top_tokens[salience_field] = []
  
          for cluster_id in range(n_clusters):
              weight_matrix_of_cluster = weight_matrix[
                  np.asarray(cluster_ids[salience_field]) == cluster_id
              ]
  
              # If this is empty, we don't have any data points in the current
              # cluster. This may happen when `reuse_clustering` is true and we get
              # only a single example.
              if not weight_matrix_of_cluster.size:
                  continue
  
              # <float32>[vocab size]
              mean_weight_matrix = weight_matrix_of_cluster.mean(axis=0)
              num_tokens_to_display = int(
                  config.get(
                      N_TOKENS_TO_DISPLAY_KEY,
                      self.config_spec()[
                          N_TOKENS_TO_DISPLAY_KEY
                      ].default,
                  )
              )
              top_indices = np.abs(mean_weight_matrix).argsort()[::-1][
                  :num_tokens_to_display
              ]
              top_tokens = [
                  (self.vocab[i], mean_weight_matrix[i])
                  for i in top_indices
                  if mean_weight_matrix[i] != 0.0
              ]
              salience_field_to_top_tokens[salience_field].append(
                  top_tokens
              )
  
      return {
          CLUSTER_ID_KEY: cluster_ids,
          REPRESENTATION_KEY: salience_field_to_representations,
          TOP_TOKEN_KEY: salience_field_to_top_tokens,
      }
  
file_path: lit/lit_nlp/components/salience_clustering.py
imports: from collections.abc import Sequence
  from typing import Optional
  from lit_nlp.api import components
  from lit_nlp.api import dataset
  from lit_nlp.api import model
  from lit_nlp.api import types
  import numpy
  from sklearn import cluster
label_desc: Use image node_id 5 for calling the SalienceClustering obj's underlying member method code with example usage: obj.run(inputs, model, dataset, model_outputs, config) and returns: dict, None