@Saeid Hedayati​ :

To store the pickle files along with the MLflow model, you can include them as artifacts when logging the model. You can modify your training script as follows:

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 artifacts
    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)
    
    # Log the PyFunc model with artifacts
    mlflow.pyfunc.log_model(pyfunc_model, "custom_model", artifacts={
        "tokenizer": tokenizer_artifact_path,
        "label_encoder": label_encoder_artifact_path
    })
    
    # 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")

In the above code, the artifacts (i.e., the pickle files) are logged along with the PyFunc model using the

mlflow.pyfunc.log_model() method. The artifacts are specified as a dictionary where the keys are the names of the artifacts and the values are the paths to the artifact files.

To load the model and the artifacts in another workspace, you can use the following code:

import mlflow.pyfunc
import joblib
 
# Load the model from the MLflow Model Registry
mlflow.set_registry_uri("my model_registry_uri")
model = mlflow.pyfunc.load_model("model_uri")
 
# Load the artifacts
tokenizer_path = model.metadata['signature_def']['serving_default']['inputs']['tokenizer'].string_value
label_encoder_path = model.metadata['signature_def']['serving_default']['inputs']['label_encoder'].string_value
tokenizer = joblib.load(tokenizer_path)
label_encoder = joblib.load(label_encoder_path)
 
# Get the PyFunc model and predict on new data
unwrapped_model = model._get_unwrapped_model()
y_pred = unwrapped_model.predict(input_data)

In the above code, we load the model and then extract the paths to the artifacts from the model metadata. We then load the artifacts using joblib.load() and use them to predict on new data.