class_docstr: Salience map from Integrated Gradients.
  
  Integrated Gradients is an attribution method originally proposed in
  Sundararajan et al. (https://arxiv.org/abs/1703.01365), which attributes an
  importance value for each input feature based on the gradients of the model
  output with respect to the input. The feature attribution values are
  calculated by taking the integral of gradients along a straight path from a
  baseline to the input being analyzed. The original implementation can be
  found at: https://github.com/ankurtaly/Integrated-Gradients/blob/master/
  BertModel/bert_model_utils.py
  
  This component requires that the following fields in the model spec. Field
  names like `embs` are placeholders; you can call them whatever you like,
  and as with other LIT components multiple segments are supported.
    Output:
      - TokenEmbeddings (`embs`) to return the input embeddings
      - TokenGradients (`grads`) to return gradients w.r.t. `embs`
      - A label field (`target`) to return the label that `grads`
        was computed for. This is usually a CategoryLabel, but can be anything
        since it will just be fed back into the model.
    Input
      - TokenEmbeddings (`embs`) to accept the modified input embeddings
      - A label field to (`target`) to pin the gradient target to the same
        label for all integral steps, since the argmax prediction may change.
class_docstr_tok: ["Salience","map","from","Integrated","Gradients",".","Integrated","Gradients","is","an","attribution","method","originally","proposed","in","Sundararajan","et","al",".","(","https",":","\/\/arxiv.org\/abs\/1703.01365",")",",","which","attributes","an","importance","value","for","each","input","feature","based","on","the","gradients","of","the","model","output","with","respect","to","the","input",".","The","feature","attribution","values","are","calculated","by","taking","the","integral","of","gradients","along","a","straight","path","from","a","baseline","to","the","input","being","analyzed",".","The","original","implementation","can","be","found","at",":","https",":","\/\/github.com\/ankurtaly\/Integrated-Gradients\/blob\/master\/","BertModel\/bert_model_utils.py","This","component","requires","that","the","following","fields","in","the","model","spec",".","Field","names","like","`","embs","`","are","placeholders",";","you","can","call","them","whatever","you","like",",","and","as","with","other","LIT","components","multiple","segments","are","supported",".","Output",":","-","TokenEmbeddings","(","`","embs","`",")","to","return","the","input","embeddings","-","TokenGradients","(","`","grads","`",")","to","return","gradients","w.r.t",".","`","embs","`","-","A","label","field","(","`","target","`",")","to","return","the","label","that","`","grads","`","was","computed","for",".","This","is","usually","a","CategoryLabel",",","but","can","be","anything","since","it","will","just","be","fed","back","into","the","model",".","Input","-","TokenEmbeddings","(","`","embs","`",")","to","accept","the","modified","input","embeddings","-","A","label","field","to","(","`","target","`",")","to","pin","the","gradient","target","to","the","same","label","for","all","integral","steps",",","since","the","argmax","prediction","may","change","."]
code: def get_salience_result(
      self,
      model_input: JsonDict,
      model: lit_model.Model,
      interpolation_steps: int,
      normalize: bool,
      class_to_explain: str,
      model_output: JsonDict,
      grad_fields: list[str],
  ):
      result = {}
  
      output_spec = model.output_spec()
      # We ensure that the embedding and gradient class fields are present in the
      # model's input spec in find_fields().
      embeddings_fields = [
          cast(types.TokenGradients, output_spec[grad_field]).grad_for
          for grad_field in grad_fields
      ]
  
      # The gradient class input is used to specify the target class of the
      # gradient calculation (if unspecified, this option defaults to the argmax,
      # which could flip between interpolated inputs).
      # If class_to_explain is emptystring, then explain the argmax class.
      grad_class_key = cast(
          types.TokenGradients, output_spec[grad_fields[0]]
      ).grad_target_field_key
      if (
          class_to_explain == ""
      ):  # pylint: disable=g-explicit-bool-comparison
          if grad_class_key in model_output:
              grad_class = model_output[grad_class_key]
          else:
              # If this is not present, should error because we can't infer
              # what class to use.
              grad_class = model_input[grad_class_key]
      else:
          grad_class = class_to_explain
  
      interpolated_inputs = {}
      all_embeddings = []
      all_baselines = []
      for embed_field in embeddings_fields:
          # <float32>[num_tokens, emb_size]
          embeddings = np.array(model_output[embed_field])
          all_embeddings.append(embeddings)
  
          # Starts with baseline of zeros. <float32>[num_tokens, emb_size]
          baseline = self.get_baseline(embeddings)
          all_baselines.append(baseline)
  
          # Get interpolated inputs from baseline to original embedding.
          # <float32>[interpolation_steps, num_tokens, emb_size]
          interpolated_inputs[
              embed_field
          ] = self.get_interpolated_inputs(
              baseline, embeddings, interpolation_steps
          )
  
      # Create model inputs and populate embedding field(s).
      inputs_with_embeds = []
      for i in range(interpolation_steps):
          # Interpolates embeddings for all inputs simultaneously.
          # Each entry is <float32>[num_tokens, emb_size]
          updates = {
              k: interpolated_inputs[k][i] for k in embeddings_fields
          }
          updates[grad_class_key] = grad_class
          input_copy = utils.make_modified_input(
              model_input, updates, "IG"
          )
          inputs_with_embeds.append(input_copy)
  
      embed_outputs = model.predict(inputs_with_embeds)
  
      # Create list with concatenated gradients for each interpolate input.
      gradients = []
      for o in embed_outputs:
          # <float32>[total_num_tokens, emb_size]
          interp_gradients = np.concatenate(
              [o[field] for field in grad_fields]
          )
          gradients.append(interp_gradients)
      # <float32>[interpolation_steps, total_num_tokens, emb_size]
      path_gradients = np.stack(gradients, axis=0)
  
      # Calculate integral
      # <float32>[total_num_tokens, emb_size]
      integral = self.estimate_integral(path_gradients)
  
      # <float32>[total_num_tokens, emb_size]
      concat_embeddings = np.concatenate(all_embeddings)
  
      # <float32>[total_num_tokens, emb_size]
      concat_baseline = np.concatenate(all_baselines)
  
      # <float32>[total_num_tokens, emb_size]
      integrated_gradients = integral * (
          np.array(concat_embeddings) - np.array(concat_baseline)
      )
      # Dot product of integral values and (embeddings - baseline).
      # <float32>[total_num_tokens]
      attributions = np.sum(integrated_gradients, axis=-1)
  
      # <float32>[total_num_tokens]
      scores = (
          citrus_utils.normalize_scores(attributions)
          if normalize
          else attributions
      )
  
      for grad_field in grad_fields:
          # Format as salience map result.
          token_field = cast(
              types.TokenGradients, output_spec[grad_field]
          ).align
          tokens = model_output[token_field]
  
          # Only use the scores that correspond to the tokens in this grad_field.
          # The gradients for all input embeddings were concatenated in the order
          # of the grad fields, so they can be sliced out in the same order.
          sliced_scores = scores[
              : len(tokens)
          ]  # <float32>[num_tokens in field]
          scores = scores[
              len(tokens) :
          ]  # <float32>[num_remaining_tokens]
  
          assert len(tokens) == len(sliced_scores)
          result[grad_field] = dtypes.TokenSalience(
              tokens, sliced_scores
          )
      return result
  
imports: from typing import cast, Optional
  from absl import logging
  from lit_nlp.api import components
  from lit_nlp.api import dataset
  from lit_nlp.api import dtypes
  from lit_nlp.api import model
  from lit_nlp.api import types
  from lit_nlp.components.citrus import utils
  from lit_nlp.lib import utils
  import numpy
label_desc: Use image node_id 6 for calling the IntegratedGradients obj's underlying member method code with example usage: obj.get_salience_result(model_input, model, interpolation_steps, normalize, class_to_explain, model_output, grad_fields) and returns: result