added data_split

udacity · Apr 2, 2024 · 3c9964b · 3c9964b
1 parent 531a02c
commit 3c9964b
Showing 1 changed file with 12 additions and 4 deletions.
diff --git a/main.py b/main.py
@@ -80,10 +80,18 @@ def go(config: DictConfig):
 
 
         if "data_split" in active_steps:
-            ##################
-            # Implement here #
-            ##################
-            pass
+            # Download file and load in W&B
+            _ = mlflow.run(
+                f"{config['main']['components_repository']}/train_val_test_split",
+                "main",
+                version='main',
+                parameters={
+                    "input": "clean_sample.csv:latest",
+                    "test_size": config["modeling"]["test_size"],
+                    "random_seed": config["modeling"]["random_seed"],
+                    "stratify_by": config["modeling"]["stratify_by"]
+                },
+            )
 
         if "train_random_forest" in active_steps: