code: def get_class_to_explain(
      model: Any,
      pred_key: str,
      example: JsonDict,
      provided_class_to_explain: Optional[int] = None,
  ) -> Optional[int]:
      """Return the class index to explain.
  
      The provided class index can be None, in which case this method determines
      which class to explain based on the class with the highest prediction score
      for the provided example.
  
      Args:
        model: The model to run.
        pred_key: The key to the model's output field to explain.
        example: The example to explain.
        provided_class_to_explain: The class index provided to this explainer, or
          None to use the argmax prediction.
  
      Returns:
        The true class index to explain, in the range [0, class vocab length).
      """
      pred_spec = model.output_spec()[pred_key]
      # If provided_class_to_explain is -1, then explain the argmax class, for both
      # multiclass and sparse multilabel tasks.
      if provided_class_to_explain is None and isinstance(
          pred_spec,
          (types.MulticlassPreds, types.SparseMultilabelPreds),
      ):
          pred = list(model.predict([example]))[0][pred_key]
          if isinstance(pred_spec, types.MulticlassPreds):
              return np.argmax(pred)
          else:
              # For sparse multi-label, sort class/score tuples to find the
              # highest-scoring class and get its vocab index.
              pred.sort(key=lambda elem: elem[1], reverse=True)
              class_name_to_explain = pred[0][0]
              return pred_spec.vocab.index(class_name_to_explain)
      else:
          return provided_class_to_explain
  
imports: from collections.abc import Iterable
  import functools
  from typing import Any, 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 lime
  from lit_nlp.components.citrus import utils
  from lit_nlp.lib import utils
  import numpy
label_desc: Use image node_id 2 for calling a global function with example usage: get_class_to_explain(model, pred_key, example, provided_class_to_explain) and returns: provided_class_to_explain, np, pred_spec