# Copyright 2022 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """This file is to serve BQML model trained with TRANSFORM.""" import collections import os from typing import Any, Optional from bigquery_ml_utils.inference.xgboost_predictor import predictor as bqml_xgboost_predictor import tensorflow as tf from bigquery_ml_utils.tensorflow_ops.load_module import load_module gen_date_ops = load_module("_date_ops.so") gen_datetime_ops = load_module("_datetime_ops.so") gen_time_ops = load_module("_time_ops.so") gen_timestamp_ops = load_module("_timestamp_ops.so") class Predictor: """Class to feed input data into BQML model trained with TRANSFORM. It performs both preprocessing and postprocessing on the input and output. """ def __init__( self, transform_savedmodel, model_tensorflow=None, model_xgboost=None ): """Initializes a Predictor to serve BQML models trained with TRANSFORM. Args: transform_savedmodel: SavedModel pb of the TRANSFORM. model_tensorflow: BQML model in tensorflow savedmodel format. model_xgboost: BQML model in booster format. Returns: A 'Predictor' instance. """ self._transform_savedmodel = transform_savedmodel self._model_tensorflow = model_tensorflow self._model_xgboost = model_xgboost # Number of input in the Predict call. self._num_input = 0 def _get_transform_result(self, raw_input): """Gets the TRANSFORM result from the raw input data. Args: raw_input: Raw input. Returns: TRANSFORM result. """ input_dict = collections.defaultdict(list) for row in raw_input: for key, value in row.items(): input_dict[key].append(value) infer = self._transform_savedmodel.signatures['serving_default'] input_signature = infer.structured_input_signature[1] transform_input = dict() for key, value in input_dict.items(): if key not in input_signature: raise ValueError(f'"{key}" in not an input of the TRANSFORM.') transform_input[key] = tf.constant(value, input_signature[key].dtype) return infer(**transform_input) def _convert_transform_result(self, transform_result): """Converts the TRANSFORM result to a list. Args: transform_result: TRANSFORM result from the raw input. Returns: TRANSFORM results in a list. Each element in the list is the TRANSFORM result of each raw input instance. """ if len(transform_result) == 1: for value in transform_result.values(): return value.numpy().tolist() # Initialize the output as a list of dict. output = [{} for _ in range(self._num_input)] # Convert the transform_result in batch representation to a list, in which # each element is the transform result of each raw input instance. for key, value in transform_result.items(): if value.dtype == tf.string: batch_result = value.numpy().astype(str).tolist() else: batch_result = value.numpy().tolist() for i, value in enumerate(batch_result): if isinstance(value, list): output[i][key] = value.copy() else: output[i][key] = value return output def _get_tf_model_result(self, transform_result): """Gets the model result from the TRANSFORM result for tensorflow models. Args: transform_result: TRANSFORM result from the raw input. Returns: Model prediction results in a list. Each element in the list is the prediction result of each raw input instance. """ model_input = dict() for key, value in transform_result.items(): model_input[key] = value # BQML trained model uses float64 as the input dtype for all numerical # features. if value.dtype == tf.int64: model_input[key] = tf.cast(value, tf.float64) infer = None if 'predict' in self._model_tensorflow.signatures: # By default, BQML DNN model is exported with 'predict' signature. infer = self._model_tensorflow.signatures['predict'] else: infer = self._model_tensorflow.signatures['serving_default'] inference_result = infer(**model_input) # If the inference_result of the model contains only one named tensor, we # omit the name. This aligns with the TF serving response. if len(inference_result) == 1: for value in inference_result.values(): # TODO(b/253233131): Support array as the output. return value.numpy().tolist() # Initialize the output as a list of dict. output = [{} for i in range(self._num_input)] # Convert the inference_result in batch representation to a list, in which # each element is the prediction result of each raw input instance. for key, value in inference_result.items(): # TODO(b/253233131): Support array as the output. batch_result = list() # The numpy() of string tensor has type BYTE, which needs to be converted # to string. if value.dtype == tf.string: batch_result = value.numpy().astype(str).tolist() else: batch_result = value.numpy().tolist() for i, value in enumerate(batch_result): if isinstance(value, list): output[i][key] = value.copy() else: output[i][key] = value return output def _get_xgboost_model_result(self, transform_result): """Gets the model result from the TRANSFORM result for xgboost models. Args: transform_result: TRANSFORM result from the raw input. Returns: Model prediction results in a list. Each element in the list is the prediction result of each raw input instance. """ xgb_model_input = [{} for i in range(self._num_input)] # Convert the transform_result in batch representation to a list, in which # each element is the prediction result of each raw input instance. for key, value in transform_result.items(): # TODO(b/253233131): Support array as the output. batch_result = list() # The numpy() of string tensor has type BYTE, which needs to be converted # to string. if value.dtype == tf.string: batch_result = value.numpy().astype(str).tolist() else: batch_result = value.numpy().tolist() for i, value in enumerate(batch_result): if isinstance(value, list): xgb_model_input[i][key] = value.copy() else: xgb_model_input[i][key] = value return self._model_xgboost.predict(xgb_model_input) def _get_model_result(self, transform_result): """Gets the model result from the TRANSFORM result. Args: transform_result: TRANSFORM result from the raw input. Returns: Model prediction results in a list. Each element in the list is the prediction result of each raw input instance. """ if self._model_tensorflow is not None: return self._get_tf_model_result(transform_result) if self._model_xgboost is not None: return self._get_xgboost_model_result(transform_result) return self._convert_transform_result(transform_result) def predict(self, raw_input, **kwargs): """Performs prediction. Args: raw_input: A list of prediction input instances. **kwargs: A dictionary of keyword args provided as additional fields on the predict request body. Returns: A dict containing the prediction results. """ del kwargs self._num_input = len(raw_input) return self._get_model_result(self._get_transform_result(raw_input)) @classmethod def from_path(cls, model_dir): """Creates an instance of Predictor using the given path. Args: model_dir: The local directory that contains the BQML model trained with TRANSFORM. Returns: An instance of 'Predictor'. """ if not tf.io.gfile.exists(os.path.join(model_dir, 'transform')): raise ValueError('TRANSFORM subdirectory is not found in the given path.') transform_savedmodel = tf.saved_model.load( os.path.join(model_dir, 'transform') ) if tf.io.gfile.exists(os.path.join(model_dir, 'saved_model.pb')): model_tensorflow = tf.saved_model.load(model_dir) return cls(transform_savedmodel, model_tensorflow=model_tensorflow) if tf.io.gfile.exists(os.path.join(model_dir, 'model.bst')): model_xgboost = bqml_xgboost_predictor.Predictor.from_path(model_dir) return cls(transform_savedmodel, model_xgboost=model_xgboost) return cls(transform_savedmodel)