Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
kedro mlflow upload
def upload_to_mlflow(kedro_annoy_dataset: KedroAnnoyIndex, idx_to_names: Dict,
item_factors: np.array, user_factors: np.array, item_biases: np.array,
user_biases: np.array, item_rank: pd.DataFrame, params: Dict):
# store temporarily the different artifacts so mlflow facilitates the logging
with tempfile.NamedTemporaryFile(prefix="idx_to_names-") as idx_to_names_file, \
tempfile.NamedTemporaryFile(prefix="item_factors_file-") as item_factors_file, \
tempfile.NamedTemporaryFile(prefix="user_factors_file-") as user_factors_file, \
tempfile.NamedTemporaryFile(prefix="item_biases_file-") as item_biases_file, \
tempfile.NamedTemporaryFile(prefix="user_biases_file-") as user_biases_file, \
tempfile.NamedTemporaryFile(prefix="params_file-") as params_file, \
tempfile.NamedTemporaryFile(prefix="item_rank_file-", mode='w') as item_rank_file:
# save in temporary files
cloudpickle.dump(idx_to_names, idx_to_names_file)
cloudpickle.dump(params, params_file)
np.save(item_factors_file, item_factors)
np.save(user_factors_file, user_factors)
np.save(item_biases_file, item_biases)
np.save(user_biases_file, user_biases)
item_rank.to_csv(item_rank_file, index=False)
# flush files to disk
idx_to_names_file.flush()
item_factors_file.flush()
user_factors_file.flush()
item_biases_file.flush()
user_biases_file.flush()
params_file.flush()
item_rank_file.flush()
artifacts = {
"annoy_index" : kedro_annoy_dataset._filepath.as_posix(),
"idx_to_names": idx_to_names_file.name,
"item_factors": item_factors_file.name,
"user_factors": user_factors_file.name,
"item_biases": item_biases_file.name,
"user_biases": user_biases_file.name,
"item_rank": item_rank_file.name,
"params": params_file.name
}
mlflow_model_logger = MlflowModelLoggerDataSet(
flavor="mlflow.pyfunc",
pyfunc_workflow="python_model",
save_args={
"artifacts": artifacts
},
)
mlflow_model_logger.save(KedroMLFlowLightFM())
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment