The unified layer for ML and AI

import torchfrom zenml import step @step(enable_cache=False)def train_model(    X: torch.Tensor, y: torch.Tensor) -> torch.nn.Module:    model = torch.nn.Sequential(        torch.nn.Linear(784, 256), torch.nn.ReLU(),        torch.nn.Linear(256, 10),    )    return model  # auto-versioned by ZenML

import tensorflow as tffrom zenml import step @step(enable_cache=False)def train_classifier(    X_train: tf.Tensor, y_train: tf.Tensor) -> tf.keras.Model:    model = tf.keras.Sequential([        tf.keras.layers.Dense(128, activation='relu'),        tf.keras.layers.Dense(10, activation='softmax'),    ])    return model  # saved as SavedModel artifact

import pandas as pdfrom sklearn.ensemble import RandomForestClassifierfrom typing_extensions import Annotatedfrom zenml import ArtifactConfig, step @stepdef model_trainer(    dataset_trn: pd.DataFrame,) -> Annotated[RandomForestClassifier,               ArtifactConfig(is_model_artifact=True)]:    model = RandomForestClassifier()    model.fit(dataset_trn.drop('target', axis=1), dataset_trn['target'])    return model

import pandas as pdfrom sklearn.datasets import load_breast_cancerfrom typing_extensions import Annotatedfrom zenml import step @stepdef data_loader(    random_state: int,) -> Annotated[pd.DataFrame, 'dataset']:    df = load_breast_cancer(as_frame=True).frame    df.reset_index(drop=True, inplace=True)    return df  # versioned DataFrame artifact

from transformers import AutoModelForCausalLMfrom peft import get_peft_model, LoraConfigfrom zenml import step @step(enable_cache=False)def finetune_step(    base_model_name: str, datasets_dir: str) -> str:    model = AutoModelForCausalLM.from_pretrained(base_model_name)    model = get_peft_model(model, LoraConfig(r=16, lora_alpha=32))    # trainer.train() — ZenML tracks the checkpoint    return datasets_dir

import pandas as pdimport xgboost as xgbfrom zenml import step @stepdef train_xgb_model(    df_train: pd.DataFrame, label_col: str = 'target') -> xgb.Booster:    dtrain = xgb.DMatrix(        df_train.drop(columns=[label_col]), df_train[label_col]    )    return xgb.train({'max_depth': 6}, dtrain, num_boost_round=100)

import lightgbm as lgbimport pandas as pdfrom zenml import step @stepdef train_lgbm(    df_train: pd.DataFrame, label: str = 'target') -> lgb.LGBMClassifier:    clf = lgb.LGBMClassifier(n_estimators=300, learning_rate=0.05)    clf.fit(df_train.drop(columns=[label]), df_train[label])    return clf

import numpy as npfrom typing_extensions import Annotatedfrom zenml import step @stepdef normalize_features(    X_raw: np.ndarray,) -> tuple[    Annotated[np.ndarray, 'X_norm'],    Annotated[np.ndarray, 'mean'],]:    mean = X_raw.mean(axis=0)    return (X_raw - mean) / X_raw.std(axis=0), mean

import polars as plfrom typing_extensions import Annotatedfrom zenml import step @stepdef build_features(    raw_path: str,) -> Annotated[pl.DataFrame, 'features']:    return (        pl.scan_parquet(raw_path)        .filter(pl.col('value') > 0)        .collect()    )

import wandbfrom zenml import stepfrom zenml.integrations.wandb.flavors import WandbExperimentTrackerSettings @step(    experiment_tracker='wandb_tracker',    settings={'experiment_tracker.wandb':        WandbExperimentTrackerSettings(tags=['training', 'v2'])})def train_and_log(X_train, y_train) -> float:    wandb.log({'loss': 0.42, 'accuracy': 0.91})    return 0.91