Streamline Your ML Workflows

from zenml import pipeline, step
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import mean_squared_error

@step
def ingest_data():
    return pd.read_csv("data/dataset.csv")

@step
def train_model(df):
    X, y = df.drop("target", axis=1), df["target"]
    model = RandomForestRegressor(n_estimators=100)
    model.fit(X, y)
    return model

@step
def evaluate_model(model, df):
    X, y = df.drop("target", axis=1), df["target"]
    rmse = mean_squared_error(y, model.predict(X)) ** 0.5
    print(f"RMSE: {rmse}")

@pipeline
def ml_pipeline():
    df = ingest_data()
    model = train_model(df)
    evaluate_model(model, df)

ml_pipeline()

from metaflow import FlowSpec, step, IncludeFile
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import mean_squared_error

class MLFlow(FlowSpec):

    data = IncludeFile("dataset.csv")

    @step
    def start(self):
        self.df = pd.read_csv(self.data.path)
        self.next(self.train_model)

    @step
    def train_model(self):
        X, y = self.df.drop("target", axis=1), self.df["target"]
        self.model = RandomForestRegressor(n_estimators=100)
        self.model.fit(X, y)
        self.next(self.evaluate)

    @step
    def evaluate(self):
        X, y = self.df.drop("target", axis=1), self.df["target"]
        self.rmse = mean_squared_error(y, self.model.predict(X)) ** 0.5
        self.next(self.end)

    @step
    def end(self):
        print(f"RMSE: {self.rmse}")

if __name__ == "__main__":
    MLFlow()