Search code examples
pythontensorflowkerasmlflow

Automatic flavor detection with ml_flow load_model?


I am using mlflow and want to handle different flavors (e.g. sklearn , tensorflow and keras) while loading.

Actually I only find the information about the stored flavor as string in

Run.to_dictionary()['data']['tags']['mlflow.log-model.history']

output:

[{"run_id": "8ea47843f7b446828dbd9cd3a1ed2339", "artifact_path": "model", "utc_time_created": "2022-04-26 10:39:55.639791", "flavors": {"keras": {"keras_module": "tensorflow.keras", "keras_version": "2.7.0", "save_format": "tf", "data": "data", "code": null}, "python_function": {"loader_module": "mlflow.keras", "python_version": "3.8.10", "data": "data", "env": "conda.yaml"}}, "model_uuid": "3bd37bdb0aa1409aabc65f8314018642", "mlflow_version": "1.25.1"}]

Run is the mlflow.entities.Run object.

Using ast.literal_eval to transform the string into the dictionary fails.

ast.literal_eval(str(self.run.to_dictionary()['data']['tags']['mlflow.log-model.history'][1:-1]))

Solution

  • The documentation of mlflow.pyfunc says:

    The python_function model flavor serves as a default model interface for MLflow Python models. Any MLflow Python model is expected to be loadable as a python_function model.

    This means you can use mlflow.pyfunc.load_model to load all mlflow supported flavors.