...

program/metaflow/helloworld.py

@@ -117,16 +117,6 @@ def load_data(self):
 
  print(f"Loaded dataset with {len(self.data)} samples")
 
- self.next(self.split_dataset)
-
- @step
- def split_dataset(self):
- """Split the data into train, validation, and test."""
- from sklearn.model_selection import train_test_split
-
- self.df_train, temp = train_test_split(self.data, test_size=0.3)
- self.df_validation, self.df_test = train_test_split(temp, test_size=0.5)
-
  self.next(self.setup_target_transformer)
 
  @step
@@ -168,6 +158,16 @@ def setup_features_transformer(self):
  ],
  )
 
+ self.next(self.split_dataset)
+
+ @step
+ def split_dataset(self):
+ """Split the data into train, validation, and test."""
+ from sklearn.model_selection import train_test_split
+
+ self.df_train, temp = train_test_split(self.data, test_size=0.3)
+ self.df_validation, self.df_test = train_test_split(temp, test_size=0.5)
+
  self.next(self.transform_target)
 
  @step
@@ -235,7 +235,7 @@ def tune_model(self):
  validation_data=(self.X_validation, self.y_validation),
  batch_size=32,
  epochs=50,
- verbose=2,
+ verbose=,
  )
 
  tuner.results_summary()

program/metaflow/requirements.txt

@@ -0,0 +1 @@
+metaflow

program/metaflow/training.py

@@ -0,0 +1,264 @@
+from io import StringIO
+from pathlib import Path
+
+import pandas as pd
+from metaflow import S3, FlowSpec, Parameter, pypi, pypi_base, step
+
+
+def load_data_from_s3(location: str):
+ print(f"Loading dataset from location {location}")
+
+ with S3(s3root=location) as s3:
+ files = s3.get_all()
+
+ print(f"Found {len(files)} file(s) in remote location")
+
+ raw_data = [pd.read_csv(StringIO(file.text)) for file in files]
+ return pd.concat(raw_data)
+
+
+def load_data_from_file():
+ location = Path("../penguins.csv")
+ print(f"Loading dataset from location {location.as_posix()}")
+ return pd.read_csv(location)
+
+
+def build_model(nodes, learning_rate):
+ from keras import Input
+ from keras.layers import Dense
+ from keras.models import Sequential
+ from keras.optimizers import SGD
+
+ model = Sequential(
+ [
+ Input(shape=(7,)),
+ Dense(nodes, activation="relu"),
+ Dense(8, activation="relu"),
+ Dense(3, activation="softmax"),
+ ],
+ )
+
+ model.compile(
+ optimizer=SGD(learning_rate=learning_rate),
+ loss="sparse_categorical_crossentropy",
+ metrics=["accuracy"],
+ )
+
+ return model
+
+
+def build_tuner_model(hp):
+ nodes = hp.Int("nodes", 10, 20, step=2)
+
+ learning_rate = hp.Float(
+ "learning_rate",
+ 1e-3,
+ 1e-2,
+ sampling="log",
+ default=1e-2,
+ )
+
+ return build_model(nodes, learning_rate)
+
+
+@pypi_base(
+ python="3.10.14",
+ packages={
+ "python-dotenv": "1.0.1",
+ "scikit-learn": "1.4.1.post1",
+ "pandas": "2.2.1",
+ "numpy": "1.26.4",
+ },
+)
+class TrainingFlow(FlowSpec):
+ debug = Parameter(
+ "debug",
+ help="Whether we are debugging the flow in a local environment",
+ default=False,
+ )
+
+ dataset_location = Parameter(
+ "dataset_location",
+ help="Location to the initial dataset",
+ default="metaflow/data/",
+ )
+
+ @step
+ def start(self):
+ self.my_var = "hello world"
+
+ self.next(self.load_data)
+
+ @pypi(packages={"boto3": "1.34.70"})
+ @step
+ def load_data(self):
+ """Load the dataset in memory.
+
+ This function reads every CSV file available and
+ concatenates them into a single dataframe.
+ """
+ import os
+
+ if self.debug:
+ df = load_data_from_file()
+ else:
+ location = f"s3://{os.environ['BUCKET']}/{self.dataset_location}"
+
+ df = load_data_from_s3(location)
+
+ # Shuffle the data
+ self.data = df.sample(frac=1, random_state=42)
+
+ print(f"Loaded dataset with {len(self.data)} samples")
+
+ self.next(self.setup_target_transformer)
+
+ @step
+ def setup_target_transformer(self):
+ from sklearn.compose import ColumnTransformer
+ from sklearn.preprocessing import OrdinalEncoder
+
+ self.target_transformer = ColumnTransformer(
+ transformers=[("species", OrdinalEncoder(), [0])],
+ )
+
+ self.next(self.setup_features_transformer)
+
+ @step
+ def setup_features_transformer(self):
+ from sklearn.compose import ColumnTransformer, make_column_selector
+ from sklearn.impute import SimpleImputer
+ from sklearn.pipeline import make_pipeline
+ from sklearn.preprocessing import OneHotEncoder, StandardScaler
+
+ numeric_transformer = make_pipeline(
+ SimpleImputer(strategy="mean"),
+ StandardScaler(),
+ )
+
+ categorical_transformer = make_pipeline(
+ SimpleImputer(strategy="most_frequent"),
+ OneHotEncoder(),
+ )
+
+ self.features_transformer = ColumnTransformer(
+ transformers=[
+ (
+ "numeric",
+ numeric_transformer,
+ make_column_selector(dtype_exclude="object"),
+ ),
+ ("categorical", categorical_transformer, ["island"]),
+ ],
+ )
+
+ self.next(self.split_dataset)
+
+ @step
+ def split_dataset(self):
+ """Split the data into train, validation, and test."""
+ from sklearn.model_selection import KFold
+
+ self.target = self.data["species"]
+ self.features = self.data.drop("species", axis=1)
+
+ kfold = KFold(n_splits=5, shuffle=True)
+ self.folds = list(enumerate(kfold.split(self.target, self.features)))
+
+ self.next(self.transform_target, foreach="folds")
+
+ @step
+ def transform_target(self):
+ import numpy as np
+
+ self.fold, (self.train_indices, self.test_indices) = self.input
+
+ self.y_train = np.squeeze(
+ self.target_transformer.fit_transform(
+ np.array(self.target.iloc[self.train_indices]).reshape(-1, 1),
+ ),
+ )
+ self.y_test = np.squeeze(
+ self.target_transformer.transform(
+ np.array(self.target.iloc[self.test_indices]).reshape(-1, 1),
+ ),
+ )
+
+ self.next(self.transform_features)
+
+ @step
+ def transform_features(self):
+ self.x_train = self.features_transformer.fit_transform(
+ self.features.iloc[self.train_indices],
+ )
+ self.x_test = self.features_transformer.transform(
+ self.features.iloc[self.test_indices],
+ )
+
+ self.next(self.train_model)
+
+ @pypi(
+ packages={
+ "keras": "3.3.0",
+ "jax[cpu]": "0.4.26",
+ "packaging": "24.0",
+ },
+ )
+ @step
+ def train_model(self):
+ print(f"Training fold {self.fold}...")
+
+ self.model = build_model(10, 0.01)
+
+ self.model.fit(
+ self.x_train,
+ self.y_train,
+ epochs=50,
+ batch_size=32,
+ verbose=2,
+ )
+
+ self.next(self.evaluate_model)
+
+ @pypi(
+ packages={
+ "keras": "3.3.0",
+ "jax[cpu]": "0.4.26",
+ "packaging": "24.0",
+ },
+ )
+ @step
+ def evaluate_model(self):
+ print(f"Evaluating fold {self.fold}...")
+
+ self.loss, self.accuracy = self.model.evaluate(
+ self.x_test,
+ self.y_test,
+ verbose=2,
+ )
+
+ print(f"Fold {self.fold} - loss: {self.loss} - accuracy: {self.accuracy}")
+ self.next(self.reduce)
+
+ @step
+ def reduce(self, inputs):
+ import numpy as np
+
+ accuracies = [i.accuracy for i in inputs]
+ accuracy = np.mean(accuracies)
+ accuracy_std = np.std(accuracies)
+
+ print(f"Accuracy: {accuracy} +-{accuracy_std}")
+
+ self.next(self.end)
+
+ @step
+ def end(self):
+ print("the end")
+
+
+if __name__ == "__main__":
+ from dotenv import load_dotenv
+
+ load_dotenv()
+ TrainingFlow()