From 7db1c4dc0040071f6b0759665b5a6a6d02ed9317 Mon Sep 17 00:00:00 2001
From: Mark Aron Szulyovszky <mark.szulyovszky@gmail.com>
Date: Sat, 19 Feb 2022 15:06:05 +0100
Subject: [PATCH] feat(Config): added flag to save models

---
 config/presets.py | 1 +
 config/types.py   | 2 ++
 run_pipeline.py   | 3 ++-
 3 files changed, 5 insertions(+), 1 deletion(-)

diff --git a/config/presets.py b/config/presets.py
index b4ff76a..cd2e3b3 100644
--- a/config/presets.py
+++ b/config/presets.py
@@ -32,4 +32,5 @@ def get_default_config() -> RawConfig:
         event_filter="cusum_fixed",
         labeling="two_class",
         forecasting_horizon=50,
+        save_models=False,
     )
diff --git a/config/types.py b/config/types.py
index 3dabda3..106f752 100644
--- a/config/types.py
+++ b/config/types.py
@@ -28,6 +28,7 @@ class RawConfig(BaseModel):
     event_filter: Literal["none", "cusum_vol", "cusum_fixed"]
     labeling: Literal["two_class", "three_class_balanced", "three_class_imbalanced"]
     forecasting_horizon: int
+    save_models: bool
 
     directional_models: list[str]
     meta_models: list[str]
@@ -50,6 +51,7 @@ class Config:
     labeling: EventLabeller
     forecasting_horizon: int
     no_of_classes: Literal["two", "three-balanced", "three-imbalanced"]
+    save_models: bool
 
     mode: Literal["training", "inference"]
 
diff --git a/run_pipeline.py b/run_pipeline.py
index 60bb5ad..123120f 100644
--- a/run_pipeline.py
+++ b/run_pipeline.py
@@ -35,7 +35,8 @@ def run_pipeline(
         wandb,
         sweep,
     )
-    save_models(pipeline_outcome, config)
+    if config.save_models:
+        save_models(pipeline_outcome, config)
     return pipeline_outcome, config