trying logging or saving model using tensorflow in mlflow

Viewed 355

I'm trying to log and/or save a model in mlflow and I get the error 'can't pickle repeated message fields, convert to list first'. I want to use the universal-sentence-encoder-large/5 model from tensorflow_hub, which can be loaded using this code:

import tensorflow_hub as hub
module_url = "https://tfhub.dev/google/universal-sentence-encoder-large/5"
model = hub.load(module_url)

The meaning of model is that it takes a string as input and finds the most similar question from the QADict below and returns an answer. You can see the model code here:

import tensorflow as tf
import numpy as np
import mlflow

QADict = {'What is your full name?': 'My name is King George.',
 'Have you taken your medication?': "I don't think so."

class convEngine(mlflow.pyfunc.PythonModel):
   
    def __init__(self, QADict, model):
        
        self.model = model
        self.Qlist = list(QADict.keys())
        self.Qanswers = self.embed(self.Qlist)
        self.QADict = QADict
        
    def embed(self, input):
      return self.model(input)
    
    def predict(self,model_input_string):
      
        qasked = self.embed([model_input_string])
        corr = np.inner(qasked, self.Qanswers)
        index_max = np.argmax(corr)
     

        return self.QADict[self.Qlist[index_max]]

When I try to log or save the model in mlflow by using the code below I get 'can't pickle repeated message fields, convert to list first'

model_path = "path/to/save/model"
convModel = convEngine(QADict, model)
mlflow.pyfunc.save_model(path=model_path , python_model=convModel)

I have also tried to log and/or save the universal-sentence-encoder-large/5 using mlflow.tensorflow.log_model() thinking on calling this model in convEngine model but without luck.

I hope some of you can help me.

0 Answers
Related