code: def infer_main(partition, config):
      import tensorflow as tf
      from tensorflow.core.protobuf import rewriter_config_pb2
      from utils import calc_auc
  
      seed = config["seed"]
      np.random.seed(seed)
      random.seed(seed)
      if tf.__version__[0] == "1":
          tf.compat.v1.set_random_seed(seed)
      elif tf.__version__[0] == "2":
          tf.random.set_seed(seed)
  
      init_start = time.time()
      os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
      os.environ["KMP_BLOCKTIME"] = "1"
      os.environ["OMP_NUM_THREADS"] = str(config["cores_per_instance"])
      os.environ["KMP_AFFINITY"] = config["kmp_affinity"]
      os.environ["KMP_SETTINGS"] = "1"
      if config["AMX"]:
          os.environ["DNNL_MAX_CPU_ISA"] = "AVX512_CORE_AMX"
      if config["dnnl_verbose"]:
          os.environ["MKL_VERBOSE"] = "1"
          os.environ["DNNL_VERBOSE"] = "1"
      init_end = time.time()
      init_time = init_end - init_start
  
      model_restore_start = time.time()
      with tf.io.gfile.GFile(config["model_path"], "rb") as f:
          graph_def = tf.compat.v1.GraphDef()
          graph_def.ParseFromString(f.read())
      with tf.Graph().as_default() as graph:
          tf.import_graph_def(graph_def, name="")
  
      input_layers = [
          "Inputs/uid_batch_ph",
          "Inputs/mid_batch_ph",
          "Inputs/cat_batch_ph",
          "Inputs/mid_his_batch_ph",
          "Inputs/cat_his_batch_ph",
          "Inputs/mask",
          "Inputs/target_ph",
      ]
  
      if config["graph_type"] == "dynamic":
          input_layers.append("Inputs/seq_len_ph")
  
      input_tensor = [
          graph.get_tensor_by_name(x + ":0") for x in input_layers
      ]
      output_layers = ["dien/fcn/add_6", "dien/fcn/Metrics/Mean_1"]
      output_tensor = [
          graph.get_tensor_by_name(x + ":0") for x in output_layers
      ]
  
      dtype = config["data_type"]
      if dtype == "bfloat16" or dtype == "int8":
          graph_options = tf.compat.v1.GraphOptions(
              rewrite_options=rewriter_config_pb2.RewriterConfig(
                  remapping=rewriter_config_pb2.RewriterConfig.AGGRESSIVE,
                  auto_mixed_precision_mkl=rewriter_config_pb2.RewriterConfig.ON,
              )
          )
      elif dtype == "fp32":
          graph_options = tf.compat.v1.GraphOptions(
              rewrite_options=rewriter_config_pb2.RewriterConfig(
                  remapping=rewriter_config_pb2.RewriterConfig.AGGRESSIVE
              )
          )
      else:
          invalidInputError(False, f"Unsupported data type: {dtype}")
  
      session_config = tf.compat.v1.ConfigProto(
          graph_options=graph_options
      )
      session_config.intra_op_parallelism_threads = config[
          "num_intra_threads"
      ]
      session_config.inter_op_parallelism_threads = config[
          "num_inter_threads"
      ]
      sess = tf.compat.v1.Session(graph=graph, config=session_config)
      model_restore_end = time.time()
      model_restore_time = model_restore_end - model_restore_start
  
      infer_start = time.time()
      accuracy_sum = 0.0
      stored_arr = []
      eval_time = 0
      sample_freq = 9999999
      options = tf.compat.v1.RunOptions(
          trace_level=tf.compat.v1.RunOptions.FULL_TRACE
      )
      run_metadata = tf.compat.v1.RunMetadata()
  
      data_size = 0
      i = 0
      buffer = []
      while True:
          run = False
          row = next(partition, [])
          if not row:  # End of the partition
              if len(buffer) > 0:
                  run = True
              else:
                  break
          elif len(buffer) == config["batch_size"]:
              run = True
          else:
              buffer.append(row)
          if run:
              data_size += len(buffer)
              feed_data = prepare_data(buffer, config)
              i += 1
  
              start_time = time.time()
              if i % sample_freq == 0:
                  prob, acc = sess.run(
                      output_tensor,
                      options=options,
                      run_metadata=run_metadata,
                      feed_dict=dict(zip(input_tensor, feed_data)),
                  )
              else:
                  prob, acc = sess.run(
                      output_tensor,
                      feed_dict=dict(zip(input_tensor, feed_data)),
                  )
  
              end_time = time.time()
              eval_time += end_time - start_time
  
              target = feed_data[6]
  
              accuracy_sum += acc
              prob_1 = prob[:, 0].tolist()
              target_1 = target[:, 0].tolist()
  
              for p, t in zip(prob_1, target_1):
                  stored_arr.append([p, t])
              if not row:
                  break
              buffer = [row]
      if data_size > 0:
          test_auc = calc_auc(stored_arr)
          accuracy_sum = accuracy_sum / i
          infer_end = time.time()
          infer_time = infer_end - infer_start
          total_recommendations = data_size
          # total_recommendations = i * config["batch_size"]
          thpt_forward_pass = float(i * config["batch_size"]) / float(
              eval_time
          )
          return [
              [
                  test_auc,
                  accuracy_sum,
                  total_recommendations,
                  (
                      init_time,
                      model_restore_time,
                      infer_time,
                      thpt_forward_pass,
                  ),
              ]
          ]
      else:
          return [[0, 0, 0, 0]]
  
file_path: BigDL/python/friesian/example/dien/e2e_infer.py
imports: import os
  import time
  import yaml
  from bigdl.friesian.utils import SafePickle
  import random
  import numpy
  from argparse import ArgumentParser
  from bigdl.orca import init_orca_context, stop_orca_context
  from bigdl.friesian.feature import FeatureTable, StringIndex
  from bigdl.dllib.utils.log4Error import *
label_desc: Use image node_id 2 for calling a global function with example usage: infer_main(partition, config) and returns: list, list