code: def _parse_prediction_log_to_tensor_value(
      prediction_log: prediction_log_pb2.PredictionLog,
  ) -> Union[np.ndarray, Dict[str, np.ndarray]]:
      """Parses the model inference values from a PredictionLog.
  
      Args:
        prediction_log: Prediction_log_pb2.PredictionLog containing inference
          results.
  
      Returns:
        Values parsed from the PredictionLog inference result. These values are
        formated in the format expected in TFMA PREDICTION_KEY Extracts value.
      """
      log_type = prediction_log.WhichOneof("log_type")
      if log_type == "classify_log":
          assert (
              len(
                  prediction_log.classify_log.response.result.classifications
              )
              == 1
          ), (
              "We expecth the number of classifications per PredictionLog to be "
              "one because TFX-BSL RunInference expects single input/output and "
              "handles batching entirely internally."
          )
          classes = np.array(
              [
                  c.label
                  for c in prediction_log.classify_log.response.result.classifications[
                      0
                  ].classes
              ],
              dtype=object,
          )
          scores = np.array(
              [
                  c.score
                  for c in prediction_log.classify_log.response.result.classifications[
                      0
                  ].classes
              ],
              dtype=np.float32,
          )
          return {"classes": classes, "scores": scores}
      elif log_type == "regress_log":
          return np.array(
              [
                  regression.value
                  for regression in prediction_log.regress_log.response.result.regressions
              ],
              dtype=float,
          )
      elif log_type == "predict_log":
          output_tensor_name_to_tensor = {}
          for (
              k,
              v,
          ) in prediction_log.predict_log.response.outputs.items():
              output_tensor_name_to_tensor[k] = np.squeeze(
                  tf.make_ndarray(v), axis=0
              )
          # If there is only one tensor (i.e. one dictionary item), we remove the
          # tensor from the dict and return it directly. Generally, TFMA will not
          # return a dictionary with a single value.
          if len(output_tensor_name_to_tensor) == 1:
              return list(output_tensor_name_to_tensor.values())[0]
          return output_tensor_name_to_tensor
      elif log_type == "multi_inference_log":
          raise NotImplementedError(
              "MultiInferenceLog processing not implemented yet."
          )
      elif log_type == "session_log":
          raise ValueError("SessionLog processing is not supported.")
      else:
          raise NotImplementedError(f"Unsupported log_type: {log_type}")
  
file_path: model-analysis/tensorflow_model_analysis/extractors/inference_base.py
imports: from typing import Dict, Optional, Sequence, Tuple, Union
  from absl import logging
  import apache_beam
  import numpy
  import tensorflow
  from tensorflow_model_analysis import constants
  from tensorflow_model_analysis.api import types
  from tensorflow_model_analysis.proto import config_pb2
  from tensorflow_model_analysis.utils import model_util
  from tensorflow_model_analysis.utils import util
  from tensorflow.python.saved_model import loader_impl
  from tensorflow_serving.apis import prediction_log_pb2
label_desc: Use image node_id 3 for calling a global function with example usage: _parse_prediction_log_to_tensor_value(prediction_log) and returns: dict, np, output_tensor_name_to_tensor, list