code: def update_prediction(
      example: dict[str, Any],
      example_output: types.JsonDict,
      output_spec: types.JsonDict,
      pred_key: str,
  ):
      """Updates prediction score and label (if classification model) in the provided example."""
      prediction = example_output[pred_key]
      example[pred_key] = prediction
      pred_spec = output_spec[pred_key]
      if isinstance(pred_spec, types.MulticlassPreds):
          # Update label
          # TODO(lit-dev): provide a general system for handling labels on
          # generated examples.
          pred_spec = cast(types.MulticlassPreds, pred_spec)
          label_key = pred_spec.parent
          label_names = pred_spec.vocab
          pred_class = np.argmax(prediction)
          example_label = label_names[pred_class]
          example[label_key] = example_label
  
file_path: lit/lit_nlp/components/cf_utils.py
imports: import re
  from typing import Any, Optional, cast
  from lit_nlp.api import types
  import numpy
label_desc: Use image node_id 1 for calling a global function with example usage: update_prediction(example, example_output, output_spec, pred_key) without return types