TensorFlow SavedModel
import pickle
from pathlib import Path
import pandas as pd
import tensorflow as tf
PACKAGE_PATH = Path(__file__).parent
OUTPUT_COLUMN = ['probability_over_50k']
class MyModel:
def __init__(self):
# Load the model
self.model = tf.keras.models.load_model(PACKAGE_PATH / 'saved_model')
def predict(self, input_df):
# Store predictions in a DataFrame
return pd.DataFrame(self.model.predict(input_df), columns=OUTPUT_COLUMN)
def get_model():
return MyModel()Last updated
Was this helpful?