code: def insert_predictions_into_extracts(
      inference_tuple: Tuple[
          types.Extracts, Dict[str, prediction_log_pb2.PredictionLog]
      ],
      output_keypath: Sequence[str] = (constants.PREDICTIONS_KEY,),
      prediction_log_keypath: Optional[Sequence[str]] = None,
  ) -> types.Extracts:
      """Inserts tensor values from PredictionLogs into the Extracts.
  
      If prediction_log_keypath is provided, then raw prediction logs from extracts
      are
      also inserted into a separate extract.
  
      Args:
        inference_tuple: This is the output of inference. It includes the key
          forwarded extracts and a dict of model name to predicition logs.
        output_keypath: A sequence of keys to be used as the path to traverse and
          insert the outputs in the extract.
        prediction_log_keypath: A sequence of keys to be used as the path to
          traverse and insert the prediction logs in the extract.
  
      Returns:
        Extracts with the PREDICTIONS_KEY populated. If prediction_log_keypath is
        provided,
        then PREDICTION_LOG_KEY is also populated with prediction logs. Note: By
        convention,
        PREDICTIONS_KEY will point to a dictionary if there are multiple
        prediction logs and a single value if there is only one prediction log.
      """
      extracts, model_names_to_prediction_logs = inference_tuple
      model_name_to_tensors = {
          name: _parse_prediction_log_to_tensor_value(log)
          for name, log in model_names_to_prediction_logs.items()
      }
      # If there is only one model (i.e. one dictionary item), we remove the model
      # output from the dict and store it directly under the PREDICTIONS_KEY. This
      # is in line with the general TFMA pattern of not storing one-item
      # dictionaries.
      if len(model_name_to_tensors) == 1:
          value = next(iter(model_name_to_tensors.values()))
      else:
          value = model_name_to_tensors
      extracts = (
          util.copy_and_set_by_keys(  # pylint: disable=protected-access
              root=extracts,
              keypath=output_keypath,
              value=value,
          )
      )
  
      # Save un-parsed prediction log if prediction_log_keypath is provided.
      # Prediction log(s) will be saved under PREDICTION_LOG_KEY either within
      # a dictionary of model names to prediction logs (if there are multiple
      # models), or by itself (if there is only one model).
      if prediction_log_keypath:
          if len(model_names_to_prediction_logs) == 1:
              value = next(
                  iter(model_names_to_prediction_logs.values())
              )
          else:
              value = model_names_to_prediction_logs
          extracts = util.copy_and_set_by_keys(  # pylint: disable=protected-access
              root=extracts,
              keypath=prediction_log_keypath,
              value=value,
          )
      return extracts
  
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 4 for calling a global function with example usage: insert_predictions_into_extracts(inference_tuple, output_keypath, prediction_log_keypath) and returns: extracts