from zenml import pipeline, step
from zenml.integrations import mlflow
@step
def preprocess_data(data):
# Preprocess the data
preprocessed_data = ...
return preprocessed_data
@step
def train_model(preprocessed_data):
# Train the model
model = ...
mlflow.log_model(model, "model")
return model
@step
def evaluate_model(model, test_data):
# Evaluate the model
metrics = ...
mlflow.log_metrics(metrics)
return metrics
@pipeline
def ml_pipeline(data, test_data):
preprocessed_data = preprocess_data(data)
model = train_model(preprocessed_data)
metrics = evaluate_model(model, test_data)
# Run the pipeline
ml_pipeline(data, test_data)
