Load a Model in code

There are a few different ways to load a ZenML Model in code:

Load the active model in a pipeline

You can also use the active model to get the model metadata, or the associated artifacts directly as described in the starter guide:

from zenml import step, pipeline, get_step_context, pipeline, Model

@pipeline(model=Model(name="my_model"))
def my_pipeline():
    ...

@step
def my_step():
    # Get model from active step context
    mv = get_step_context().model

    # Get metadata
    print(mv.run_metadata["metadata_key"].value)

    # Directly fetch an artifact that is attached to the model
    output = mv.get_artifact("my_dataset", "my_version")
    output.run_metadata["accuracy"].value

Load any model via the Client

Alternatively, you can use the Client:

from zenml import step
from zenml.client import Client
from zenml.enums import ModelStages

@step
def model_evaluator_step()
    ...
    # Get staging model version 
    try:
        staging_zenml_model = Client().get_model_version(
            model_name_or_id="<INSERT_MODEL_NAME>",
            model_version_name_or_number_or_id=ModelStages.STAGING,
        )
    except KeyError:
        staging_zenml_model = None
    ...