Custom mlflow pyfunc wrapper for Keras models

mlflow-pyfunc-wrapper.py

import mlflow.pyfunc
import mlflow.keras

class KerasWrapper(mlflow.pyfunc.PythonModel):
  
  def __init__(self, keras_model_name):         
    self.keras_model_name = keras_model_name
    
  def load_context(self, context):
    
    self.keras_model = mlflow.keras.load_model(model_uri=context.artifacts[self.keras_model_name], compile=False)
  
  def predict(self, context, input_data):
    import pandas as pd
    import numpy as np

    input_data_padded = np.stack(input_data["0"].apply(np.stack, axis=0), axis=0)
    scores = self.keras_model.predict(input_data_padded)
    predicted_class = np.apply_along_axis(np.argmax, 1, scores)
    scores.sort(axis=1)
    marginal_confidence = np.apply_along_axis(lambda x: x[-1] - x[-2], 1, scores)
    predicted = pd.DataFrame(
      data=dict(
        predicted_class=predicted_class.astype(dtype="float32"),
        marginal_confidence=marginal_confidence)
    )
    return predicted

I wrote this to postprocess the model's class probabilities into a predicted class and confidence.

The error you're getting is because your model isn't serialisable. That's something to do with the way keras is implemented (I don't fully understand the details) but I'm afraid it will prevent you using keras and mlflow.pyfunc together until this issue gets resolved.