How to use mlflow to deploy model that requires tensorflow_text for bert on local machine? - tensorflow

Question

I recently use mlflow 1.29.0 to track my model training, I use BERT for text embedding which need to import tensorflow_text to register op before training, here is an example:
import tensorflow_hub as hub
import tensorflow_text as text
def create_model():
text_input = tf.keras.layers.Input(shape = (),dtype = tf.string, name = 'text_input')
preprocessed_text = preprocess_model(text_input)
encoder_text = encoder_model(preprocessed_text)['pooled_output']
text_output = tf.keras.layers.Dropout(0.1,name = 'dropout1')(encoder_text)
text_output = tf.keras.layers.Dense(units = 400, activation = tf.keras.activations.sigmoid, name = 'text_dense1')(text_output)
text_output = tf.keras.layers.Dropout(0.1,name = 'dropout2')(text_output)
final_output = tf.keras.layers.Dense(units = 1, activation = tf.keras.activations.sigmoid,name = 'output')(text_output)
model = tf.keras.Model(inputs = [text_input],outputs = [final_output])
return model
if __name__ == '__main__':
preprocess_path = 'https://tfhub.dev/tensorflow/bert_en_uncased_preprocess/3'
encoder_path = 'https://tfhub.dev/tensorflow/bert_en_uncased_L-24_H-1024_A-16/4'
preprocess_model = hub.KerasLayer(preprocess_path)
encoder_model = hub.KerasLayer(encoder_path)
with mlflow.start_run() as run:
model = create_model()
model.fit(...)
mlflow.keras.log_model(keras_model = model,...)
mlflow.end_run()
The code run successfully, and the mlflow ui showed everything, however, when I start to deploy the data on my local machine with the following command
mlflow sagemaker build-and-push-container
mlflow sagemaker run-local -m runs:/XXXXX/XXXX -p 4999
it showed the following error:
FileNotFoundError: Op type not registered 'CaseFoldUTF8' in binary running on mighty. Make sure the Op and Kernel are registered in the binary running in this process. Note that if you are loading a saved graph which used ops from tf.contrib, accessing (e.g.) tf.contrib.resampler should be done before importing the graph, as contrib ops are lazily registered when the module is first accessed.
I think it's because the tensorflow_text need to be import before running the model. (based on mlflow, the conda.yaml contain tensorflow-text==2.3.0)
I've meet this error several times when I train the model, I put 'import tensorflow_text as text' at top and the problem was fixed.
However I'm not quite sure how to do that when I deploy the model locally, can anyone help me with that? thank you!
I tried other command like mlflow models serve -m runs:/XXXXX/XXXX -p 4999, and the error is still there.