from mlflow import MlflowClient

# Define model name in the Model Registry
model_name = f"{catalog_name}.{schema_name}.summarizer"

# Point to Unity-Catalog registry and log/push artifact
mlflow.set_registry_uri("databricks-uc")

autolog_run = mlflow.last_active_run()
mlflow.register_model(
  "runs:/<run_id>/model", "prod.ml_team.iris_model"
)

import mlflow.pyfunc
model_version_uri = "models:/prod.ml_team.iris_model@Champion"
champion_version = mlflow.pyfunc.load_model(model_version_uri)
champion_version.predict(test_x)