docstr: Predict the input data
  :param data: predict input data.  It can be XShards or Spark DataFrame.
         If data is XShards, each partition can be a Pandas DataFrame or a dictionary of
         {'x': feature}, where feature is a numpy array or a tuple of numpy arrays.
  :param batch_size: Total batch size for all workers used for evaluation. Each worker's batch
         size would be this value divide the total number of workers. Default: 32.
  :param verbose: Prints output of one model if true.
  :param steps: Total number of steps (batches of samples) before declaring the prediction
         round finished. Ignored with the default value of None.
  :param callbacks: List of Keras compatible callbacks to apply during prediction.
  :param data_config: An optional dictionary that can be passed to data creator function.
  :param feature_cols: Feature column name(s) of data. Only used when data is a Spark
         DataFrame or an XShards of Pandas DataFrame. Default: None.
  :param output_cols: Column name(s) of the model output data. Only used when data is
         a Spark DataFrame, note the order of column name(s) should be consistent with the
         model output data. Default: None.
  :return:
docstr_tok: ["Predict","the","input","data",":","param","data",":","predict","input","data",".","It","can","be","XShards","or","Spark","DataFrame",".","If","data","is","XShards",",","each","partition","can","be","a","Pandas","DataFrame","or","a","dictionary","of","{","'","x","'",":","feature","}",",","where","feature","is","a","numpy","array","or","a","tuple","of","numpy","arrays",".",":","param","batch_size",":","Total","batch","size","for","all","workers","used","for","evaluation",".","Each","worker","'s","batch","size","would","be","this","value","divide","the","total","number","of","workers",".","Default",":","32.",":","param","verbose",":","Prints","output","of","one","model","if","true",".",":","param","steps",":","Total","number","of","steps","(","batches","of","samples",")","before","declaring","the","prediction","round","finished",".","Ignored","with","the","default","value","of","None",".",":","param","callbacks",":","List","of","Keras","compatible","callbacks","to","apply","during","prediction",".",":","param","data_config",":","An","optional","dictionary","that","can","be","passed","to","data","creator","function",".",":","param","feature_cols",":","Feature","column","name","(","s",")","of","data",".","Only","used","when","data","is","a","Spark","DataFrame","or","an","XShards","of","Pandas","DataFrame",".","Default",":","None",".",":","param","output_cols",":","Column","name","(","s",")","of","the","model","output","data",".","Only","used","when","data","is","a","Spark","DataFrame",",","note","the","order","of","column","name","(","s",")","should","be","consistent","with","the","model","output","data",".","Default",":","None",".",":","return",":"]
code: def predict(
      self,
      data: Union["SparkXShards", "SparkDataFrame"],
      batch_size: Optional[int] = 32,
      verbose: Union[str, int] = 1,
      steps: Optional[int] = None,
      callbacks: Optional[List["Callback"]] = None,
      data_config: Optional[Dict] = None,
      feature_cols: Optional[List[str]] = None,
      output_cols: Optional[List[str]] = None,
  ) -> Union["SparkXShards", "SparkDataFrame"]:
      """
      Predict the input data
      :param data: predict input data.  It can be XShards or Spark DataFrame.
             If data is XShards, each partition can be a Pandas DataFrame or a dictionary of
             {'x': feature}, where feature is a numpy array or a tuple of numpy arrays.
      :param batch_size: Total batch size for all workers used for evaluation. Each worker's batch
             size would be this value divide the total number of workers. Default: 32.
      :param verbose: Prints output of one model if true.
      :param steps: Total number of steps (batches of samples) before declaring the prediction
             round finished. Ignored with the default value of None.
      :param callbacks: List of Keras compatible callbacks to apply during prediction.
      :param data_config: An optional dictionary that can be passed to data creator function.
      :param feature_cols: Feature column name(s) of data. Only used when data is a Spark
             DataFrame or an XShards of Pandas DataFrame. Default: None.
      :param output_cols: Column name(s) of the model output data. Only used when data is
             a Spark DataFrame, note the order of column name(s) should be consistent with the
             model output data. Default: None.
      :return:
      """
      # Use the local batch size for each worker to convert to XShards
      if batch_size:
          invalidInputError(
              isinstance(batch_size, int) and batch_size > 0,
              "batch_size should be a positive integer",
          )
          local_batch_size = batch_size // self.num_workers
          if local_batch_size <= 0:
              local_batch_size = 1
      else:
          local_batch_size = None
      logger.info("Starting predict step.")
      sc = OrcaContext.get_spark_context()
      if self.model_weights:
          weights = sc.broadcast(self.model_weights)
      else:
          weights = None
  
      init_params = dict(
          model_creator=self.model_creator,
          model_load=self.load_path,
          compile_args_creator=self.compile_args_creator,
          config=self.config,
          verbose=self.verbose,
          size=self.num_workers,
          model_weights=weights,
          mode="predict",
          cluster_info=None,  # cluster_info is not needed for predict
          model_dir=self.model_dir,
          application_id=self.application_id,
          need_to_log_to_driver=self.need_to_log_to_driver,
          driver_ip=self.ip,
          driver_port=self.port,
      )
  
      if self.load_params is not None:
          init_params.update(self.load_params)
  
      params = dict(
          verbose=verbose,
          batch_size=batch_size,
          steps=steps,
          callbacks=callbacks,
          data_config=data_config,
          output_cols=output_cols,
      )
  
      def transform_func(iter, init_param, param):
          partition_data = list(iter)
          # res = combine_in_partition(partition_data)
          param["data_creator"] = make_data_creator(partition_data)
          return SparkRunner(**init_param).predict(**param)
  
      if isinstance(
          data, DataFrame
      ):  # Computation would be triggered by the user
          xshards, _ = dataframe_to_xshards(
              data,
              validation_data=None,
              feature_cols=feature_cols,
              label_cols=None,
              mode="predict",
              accept_str_col=True,
              shard_size=local_batch_size,
          )
  
          def transform_func(iter, init_param, param):
              param["data_creator"] = make_data_creator(iter)
              return SparkRunner(**init_param).predict(**param)
  
          pred_shards = SparkXShards.lazy(
              xshards.rdd.mapPartitions(
                  lambda iter: transform_func(iter, init_params, params)
              )
          )
          result = convert_predict_xshards_to_dataframe(
              data, pred_shards, output_cols
          )
      elif isinstance(
          data, SparkXShards
      ):  # Computation triggered when updating XShards
          xshards = data.to_lazy()
          if xshards._get_class_name() == "pandas.core.frame.DataFrame":
              xshards = process_xshards_of_pandas_dataframe(
                  xshards, feature_cols
              )
              pred_shards = SparkXShards.lazy(
                  xshards.rdd.mapPartitions(
                      lambda iter: transform_func(
                          iter, init_params, params
                      )
                  )
              )
              # Should add to the original SparkXShards of Pandas DataFrames
              result = add_predict_to_pd_xshards(data, pred_shards)
          else:
              pred_shards = SparkXShards.lazy(
                  xshards.rdd.mapPartitions(
                      lambda iter: transform_func(
                          iter, init_params, params
                      )
                  )
              )
              result = update_predict_xshards(data, pred_shards)
          # Uncache the original data since it is already included in the result
          data.uncache()
      else:
          invalidInputError(
              False,
              "Only XShards or Spark DataFrame are supported for predict",
          )
      return result
  
imports: import os
  import types
  import logging
  import tempfile
  import shutil
  from bigdl.dllib.utils.common import get_node_and_core_number
  from bigdl.dllib.utils.file_utils import is_local_path
  from bigdl.dllib.utils.utils import get_node_ip
  from bigdl.orca.data.file import is_file, exists, get_remote_file_to_local, get_remote_files_with_prefix_to_local, put_local_file_to_remote, put_local_files_with_prefix_to_remote
  from bigdl.orca.learn.tf2.spark_runner import SparkRunner
  from bigdl.orca.learn.utils import find_free_port, find_ip_and_free_port
  from bigdl.orca.learn.utils import maybe_dataframe_to_xshards, dataframe_to_xshards, convert_predict_xshards_to_dataframe, make_data_creator, load_model, save_model, process_xshards_of_pandas_dataframe, add_predict_to_pd_xshards, update_predict_xshards
  from bigdl.orca.learn.log_monitor import start_log_server, stop_log_server
  from bigdl.orca import OrcaContext
  from bigdl.orca.data.shard import SparkXShards
  from bigdl.dllib.utils.log4Error import invalidInputError
  from pyspark.sql.dataframe import DataFrame
  from typing import TYPE_CHECKING, Any, Dict, List, Callable, Union, Optional
label_desc: Use image node_id 5 for calling the SparkTFEstimator obj's underlying member method code with example usage: obj.predict(data, batch_size, verbose, steps, callbacks, data_config, feature_cols, output_cols) and returns: result, SparkRunner, SparkRunner