odel_artifact("trained_model")

# Make predictionspredictions = pd.Series(model.predict(data))

return predictions

However, this approach has the downside that if the step is cached, then it could lead to unexpected results. You could simply disable the cache in the above step or the corresponding pipeline. However, one other way of achieving this would be to resolve the artifact at the pipeline level:

from typing_extensions import Annotated

from zenml import get_pipeline_context, pipeline, Model

from zenml.enums import ModelStages

import pandas as pd

from sklearn.base import ClassifierMixin

@step

def predict(

model: ClassifierMixin,

data: pd.DataFrame,

) -> Annotated[pd.Series, "predictions"]:

predictions = pd.Series(model.predict(data))

return predictions

@pipeline(

model=Model(

name="iris_classifier",

# Using the production stage

version=ModelStages.PRODUCTION,

),

def do_predictions():

# model name and version are derived from pipeline context

model = get_pipeline_context().model

inference_data = load_data()

predict(

# Here, we load in the `trained_model` from a trainer step

model=model.get_model_artifact("trained_model"),

data=inference_data,

if __name__ == "__main__":

do_predictions()

Ultimately, both approaches are fine. You should decide which one to use based on your own preferences.

PreviousAssociate a pipeline with a Model

NextControlling Model versions

Last updated 19 days ago