Dear community,

I want to basically store 2 pickle files during the training and model registry with my keras model. So that when I access the model from another workspace (using mlflow.set_registery_uri()) , these models can be accessed as well. The custom mlflow model that I am using is as follow:

class KerasModel(mlflow.pyfunc.PythonModel):
 
  def __init__(self, model, tokenizer_path, label_encoder_path):
    self.model = model
    self.tokenizer_path = tokenizer_path
    self.label_encoder_path = label_encoder_path
 
  def _load_tokenizer(self):
    return joblib.load(self.tokenizer_path)
 
  def _load_label_encoder(self):
    return joblib.load(self.label_encoder_path)
 
  def predict(self, context, input_data):
    y_pred = self.model.predict(input_data)
    return y_pred

and here is my training script:

import joblib
 
import mlflow
import mlflow.keras
import mlflow.tensorflow
 
from keras.preprocessing.text import Tokenizer
from sklearn.preprocessing import LabelEncoder
 
import  keras
import tensorflow
 
# Load and preprocess data into train/test splits
X_train, y_train = get_training_data()
 
########
# do data preprocessing.....
########
 
tokenizer_artifact_path = "/dbfs/tmp/train/tokenizer.pkl"
joblib.dump(fitted_tokenizer, tokenizer_artifact_path)
 
label_encoder_artifact_path = "/dbfs/tmp/train/label_encoder.pkl"
joblib.dump(fitted_label_encoder, label_encoder_artifact_path)
 
 
with mlflow.start_run() as mlflow_run:
 
    # Fit keras model and log model
    ########
    # build keras model.....
    ########
    model, model_history = model.fit(X_train, y_train)
    mlflow.keras.log_model(model, "model")
    
 
    # log label encoder and tokenizer as artifact
    mlflow.log_artifact(tokenizer_artifact_path)
    mlflow.log_artifact(label_encoder_artifact_path)
         
    # Create a PyFunc model that uses the trained Keras model and label encoder
    pyfunc_model = KerasModel(model, tokenizer_artifact_path, label_encoder_artifact_path)
    mlflow.pyfunc.log_model("custom_model", python_model= pyfunc_model)
    
    # get mlflow artifact uri
    artifact_uri = mlflow_run.info.artifact_uri
    model_uri = artifact_uri + "/custom_model"
 
    # Register model to MLflow Model Registry if provided
    mlflow.set_registry_uri("my registery_uri")
    mlflow.register_model(model_uri, name="keras_clssification")

The problem is when I want to access this registered model from another workspace, I can load the model but not the pickle files and it throw the error that

FileNotFoundError: [Errno 2] No such file or directory: '/dbfs/tmp/train/label_encoder.pkl'

I use the following code:

mlflow.set_registery_uri("my model_registery_uri")
model = mlflow.pyfunc.load_model("model_uri")
unwrapped_model = model.unwrap_python_model()
label_encoder = unwrapped_model._load_label_encoder()
tokenizer = unwrapped_model._load_tokenizer()

It works on the same work space since the path is recognizable. But on other workspace it has not access to this. My question is how to store these two pickle files with the model so that where ever the model goes, these files goes as well?

I have checked this solution here as well, unfortunately I could not understand it completely.

If you could post your answer with code I really appreciate!

With many thanks in advance!