diff --git a/config.yaml b/config.yaml
index 3c52f30..44b3914 100644
--- a/config.yaml
+++ b/config.yaml
@@ -15,18 +15,18 @@ model:
   ]
   time_varying_known_reals: [ "time_idx" ]
   batch_size: 512
-  max_epochs: 100
+  max_epochs: 300
   gpus: 1
   learning_rate: 0.0001
-  hidden_size: 50
+  hidden_size: 100
   dropout: 0.2
-  hidden_continuous_size: 25
+  hidden_continuous_size: 50
   gradient_clip_val: 0.3
   max_prediction_length: 24
   max_encoder_length: 48
   sample: "60min"
   cutoff: 0.70
-  model_path: "/content/model/tft.pt"
+  model_path: "/content/drive/MyDrive/Colab_Notebooks/model/tft.pt"
 
 model_local:
   lags: None
@@ -42,7 +42,7 @@ model_local:
   batch_size: 1024
   max_epochs: 4
   gpus: 0
-  learning_rate: 0.0001
+  learning_rate: 0.01
   hidden_size: 8
   dropout: 0.2
   hidden_continuous_size: 4
@@ -51,4 +51,4 @@ model_local:
   max_encoder_length: 48
   sample: "60min"
   cutoff: 0.70
-  model_path: "model/tft_regressor_test.pt"
\ No newline at end of file
+  model_path: "model/temporal_fusion_transformer/tft_local.pt"
\ No newline at end of file
diff --git a/forecast.py b/forecast_all.py
similarity index 100%
rename from forecast.py
rename to forecast_all.py
diff --git a/tft_predict_colab.py b/forecast_tft.py
similarity index 61%
rename from tft_predict_colab.py
rename to forecast_tft.py
index 9cb38e5..367a400 100644
--- a/tft_predict_colab.py
+++ b/forecast_tft.py
@@ -1,58 +1,84 @@
 from os import listdir
 from os.path import isfile, join
 import warnings
+import argparse
 
 import pandas as pd
 import matplotlib.pyplot as plt
 
-import torch
-from pytorch_forecasting.metrics import SMAPE, PoissonLoss, QuantileLoss
-from pytorch_forecasting import Baseline, TemporalFusionTransformer, TimeSeriesDataSet
-from pytorch_forecasting.data import GroupNormalizer
-
 from config import load_config
 from load_data import LoadData
 
+import torch
+from pytorch_forecasting import TemporalFusionTransformer, TimeSeriesDataSet
+from pytorch_forecasting.data import GroupNormalizer
+from pytorch_forecasting.metrics import QuantileLoss
+
 warnings.filterwarnings("ignore")
 
-spec = load_config("config.yaml")
-DATA_PATH = spec["general"]["data_path"]
+parser = argparse.ArgumentParser(description="Run Training Pipeline")
+parser.add_argument(
+    "-l",
+    "--local",
+    help="Run local or in colab",
+    action='store_true',
+    default=False,
+)
+args = parser.parse_args()
+
+LOCAL = args.local
+
+if LOCAL:
+    model_key = "model_local"
+    spec = load_config("config.yaml")
+    DATA_PATH = spec["general"]["data_path"]
+    MODEL_PATH = spec[model_key]["model_path"]
+else:
+    model_key = "model"
+    import tensorflow as tf
+    import tensorboard as tb
+    tf.io.gfile = tb.compat.tensorflow_stub.io.gfile
+    spec = load_config("/content/temporal-fusion-transformer/config.yaml")
+    DATA_PATH = "/content/temporal-fusion-transformer/" + spec["general"]["data_path"]
+    MODEL_PATH = "model/temporal_fusion_transformer/tft.pt"
+
 FOLDER_LIST = spec["general"]["folder_list"]
-MODEL_PATH = spec["model"]["model_path"]
-BATCH_SIZE = spec["model"]["batch_size"]
-MAX_EPOCHS = spec["model"]["max_epochs"]
-GPUS = spec["model"]["gpus"]
-LEARNING_RATE = spec["model"]["learning_rate"]
-HIDDEN_SIZE = spec["model"]["hidden_size"]
-DROPOUT = spec["model"]["dropout"]
-HIDDEN_CONTINUOUS_SIZE = spec["model"]["hidden_continuous_size"]
-GRADIENT_CLIP_VAL = spec["model"]["gradient_clip_val"]
-
-lags = spec["model_local"]["lags"]
-sma = spec["model_local"]["sma"]
+BATCH_SIZE = spec[model_key]["batch_size"]
+MAX_EPOCHS = spec[model_key]["max_epochs"]
+GPUS = spec[model_key]["gpus"]
+LEARNING_RATE = spec[model_key]["learning_rate"]
+HIDDEN_SIZE = spec[model_key]["hidden_size"]
+DROPOUT = spec[model_key]["dropout"]
+HIDDEN_CONTINUOUS_SIZE = spec[model_key]["hidden_continuous_size"]
+GRADIENT_CLIP_VAL = spec[model_key]["gradient_clip_val"]
+
+lags = spec[model_key]["lags"]
+sma = spec[model_key]["sma"]
 sma_columns = [f"sma_{sma}" for sma in sma]
 
 if lags != "None":
     lags_columns = [f"(t-{lag})" for lag in range(lags, 0, -1)]
 
     time_varying_known_reals = (
-            spec["model_local"]["time_varying_known_reals"] +
+            spec[model_key]["time_varying_known_reals"] +
             lags_columns +
             sma_columns
     )
 if lags == "None":
     lags = None
     time_varying_known_reals = (
-            spec["model_local"]["time_varying_known_reals"] +
+            spec[model_key]["time_varying_known_reals"] +
             sma_columns
     )
 
-time_varying_known_categoricals = spec["model"]["time_varying_known_categoricals"]
+else:
+    time_varying_known_reals = spec[model_key]["time_varying_known_reals"]
 
-max_prediction_length = spec["model"]["max_prediction_length"]
-max_encoder_length = spec["model"]["max_encoder_length"]
-sample = spec["model"]["sample"]
-cutoff = spec["model"]["cutoff"]
+time_varying_known_categoricals = spec[model_key]["time_varying_known_categoricals"]
+max_prediction_length = spec[model_key]["max_prediction_length"]
+max_encoder_length = spec[model_key]["max_encoder_length"]
+sample = spec[model_key]["sample"]
+cutoff = spec[model_key]["cutoff"]
 
 train_data, test_data = LoadData(
     data_path=DATA_PATH,
@@ -67,7 +93,7 @@
     train_data,
     time_idx="time_idx",
     target="value",
-    group_ids=["dataset", "id"],
+    group_ids=["id"],
     min_encoder_length=max_encoder_length // 2,
     max_encoder_length=max_encoder_length,
     min_prediction_length=1,
@@ -78,12 +104,12 @@
     time_varying_unknown_categoricals=[],
     time_varying_unknown_reals=["value"],
     target_normalizer=GroupNormalizer(
-        groups=["dataset", "id"], transformation="softplus"
+        groups=["id"], transformation="softplus"
     ),
     add_relative_time_idx=True,
     add_target_scales=True,
     add_encoder_length=True,
-    allow_missing_timesteps=True,
+
 )
 
 model = TemporalFusionTransformer.from_dataset(
@@ -97,13 +123,12 @@
     loss=QuantileLoss(),
     log_interval=10,
     reduce_on_plateau_patience=4,
-)
+    )
 
-model.load_state_dict(torch.load("model/tft_regressor_test.pt"))
+model.load_state_dict(torch.load(MODEL_PATH))
 
 for folder in FOLDER_LIST:
 
-    print(folder)
     folder_path = f"{DATA_PATH}/{folder}"
     file_list = [
         f for f in listdir(folder_path) if isfile(join(folder_path, f))
@@ -119,7 +144,6 @@
 
         for start in range(0, 80, 1):
 
-            # start = 0
             test_data_df = df[start:(start + max_encoder_length)]
             y_obs = df[(start + max_encoder_length): (start + max_encoder_length + max_prediction_length)]
 
diff --git a/params.json b/initial_tests/params.json
similarity index 100%
rename from params.json
rename to initial_tests/params.json
diff --git a/tft_pipeline_colab.py b/initial_tests/tft_pipeline_colab.py
similarity index 99%
rename from tft_pipeline_colab.py
rename to initial_tests/tft_pipeline_colab.py
index e701b1e..ada2396 100644
--- a/tft_pipeline_colab.py
+++ b/initial_tests/tft_pipeline_colab.py
@@ -53,7 +53,6 @@
     time_varying_known_reals = (spec["model"]["time_varying_known_reals"])
 
 time_varying_known_categoricals = spec["model"]["time_varying_known_categoricals"]
-
 max_prediction_length = spec["model"]["max_prediction_length"]
 max_encoder_length = spec["model"]["max_encoder_length"]
 sample = spec["model"]["sample"]
@@ -87,11 +86,9 @@
         target_normalizer=GroupNormalizer(
             groups=["id"], transformation="softplus"
         ),
-        lags={"sma_12": [1, 24, 48]},
         add_relative_time_idx=True,
         add_target_scales=True,
         add_encoder_length=True,
-
     )
     # create validation set (predict=True) which means to predict the
     # last max_prediction_length points in time for each series
diff --git a/tft_predict_local.py b/initial_tests/tft_predict_local.py
similarity index 99%
rename from tft_predict_local.py
rename to initial_tests/tft_predict_local.py
index 8fa6916..a3e7936 100644
--- a/tft_predict_local.py
+++ b/initial_tests/tft_predict_local.py
@@ -83,7 +83,6 @@
     add_relative_time_idx=True,
     add_target_scales=True,
     add_encoder_length=True,
-    allow_missing_timesteps=True,
 )
 
 model = TemporalFusionTransformer.from_dataset(
diff --git a/model/model_cpi_colab.pt b/model/model_cpi_colab.pt
deleted file mode 100644
index 23ed913..0000000
Binary files a/model/model_cpi_colab.pt and /dev/null differ
diff --git a/model/n_beats/n_beats.pt b/model/n_beats/n_beats.pt
index 5f7e3c4..4e1d94b 100644
Binary files a/model/n_beats/n_beats.pt and b/model/n_beats/n_beats.pt differ
diff --git a/model/n_beats/training.pickle b/model/n_beats/training.pickle
index e1f097a..ed9cca1 100644
Binary files a/model/n_beats/training.pickle and b/model/n_beats/training.pickle differ
diff --git a/model/temporal_fusion_transformer/tft.pt b/model/temporal_fusion_transformer/tft.pt
index aa48d84..7d1ecc7 100644
Binary files a/model/temporal_fusion_transformer/tft.pt and b/model/temporal_fusion_transformer/tft.pt differ
diff --git a/model/temporal_fusion_transformer/tft_local.pt b/model/temporal_fusion_transformer/tft_local.pt
new file mode 100644
index 0000000..5e5f57a
Binary files /dev/null and b/model/temporal_fusion_transformer/tft_local.pt differ
diff --git a/model/tft_regressor_local.pt b/model/tft_regressor_local.pt
deleted file mode 100644
index 4293781..0000000
Binary files a/model/tft_regressor_local.pt and /dev/null differ
diff --git a/model/tft_regressor_test.pt b/model/tft_regressor_test.pt
deleted file mode 100644
index 288a07d..0000000
Binary files a/model/tft_regressor_test.pt and /dev/null differ
diff --git a/exponential_smoothing.py b/train_ets.py
similarity index 100%
rename from exponential_smoothing.py
rename to train_ets.py
diff --git a/n_beats.py b/train_nbeats.py
similarity index 98%
rename from n_beats.py
rename to train_nbeats.py
index 8ccfa10..580d03f 100644
--- a/n_beats.py
+++ b/train_nbeats.py
@@ -58,7 +58,7 @@
         time_idx="time_idx",
         target="value",
         # categorical_encoders={"id": NaNLabelEncoder().fit(train_data.id)},
-        group_ids=["id"],
+        group_ids=["dataset", "id"],
         time_varying_unknown_reals=["value"],
         max_encoder_length=max_encoder_length,
         max_prediction_length=max_prediction_length,
@@ -70,7 +70,7 @@
     validation = TimeSeriesDataSet.from_dataset(training, train_data, min_prediction_idx=training_cutoff + 1)
 
     # create dataloaders for model
-    batch_size = 512
+    batch_size = 1024
     train_dataloader = training.to_dataloader(train=True, batch_size=batch_size, num_workers=0)
     val_dataloader = validation.to_dataloader(train=False, batch_size=batch_size, num_workers=0)
 
@@ -79,7 +79,7 @@
 
     early_stop_callback = EarlyStopping(monitor="val_loss", min_delta=1e-4, patience=10, verbose=False, mode="min")
     trainer = pl.Trainer(
-        max_epochs=5,
+        max_epochs=3,
         gpus=0,
         weights_summary="top",
         gradient_clip_val=0.01,
diff --git a/sarimax.py b/train_sarimax.py
similarity index 100%
rename from sarimax.py
rename to train_sarimax.py
diff --git a/tft_pipeline_local.py b/train_tft.py
similarity index 84%
rename from tft_pipeline_local.py
rename to train_tft.py
index f11b961..d4e87d0 100644
--- a/tft_pipeline_local.py
+++ b/train_tft.py
@@ -1,4 +1,5 @@
 import warnings
+import argparse
 
 from config import load_config
 from load_data import LoadData
@@ -8,24 +9,36 @@
 from pytorch_lightning.loggers import TensorBoardLogger
 import torch
 
-from pytorch_forecasting import Baseline, TemporalFusionTransformer, TimeSeriesDataSet
+from pytorch_forecasting import TemporalFusionTransformer, TimeSeriesDataSet
 from pytorch_forecasting.data import GroupNormalizer
 from pytorch_forecasting.metrics import QuantileLoss
 
-LOCAL = True
+warnings.filterwarnings("ignore")
+
+parser = argparse.ArgumentParser(description="Run Training Pipeline")
+parser.add_argument(
+    "-l",
+    "--local",
+    help="Run local or in colab",
+    action='store_true',
+    default=False,
+)
+args = parser.parse_args()
+
+LOCAL = args.local
 
 if LOCAL:
     model_key = "model_local"
+    spec = load_config("config.yaml")
+    DATA_PATH = spec["general"]["data_path"]
 else:
     model_key = "model"
     import tensorflow as tf
     import tensorboard as tb
     tf.io.gfile = tb.compat.tensorflow_stub.io.gfile
+    spec = load_config("/content/temporal-fusion-transformer/config.yaml")
+    DATA_PATH = "/content/temporal-fusion-transformer/" + spec["general"]["data_path"]
 
-warnings.filterwarnings("ignore")
-
-spec = load_config("config.yaml")
-DATA_PATH = spec["general"]["data_path"]
 FOLDER_LIST = spec["general"]["folder_list"]
 MODEL_PATH = spec[model_key]["model_path"]
 BATCH_SIZE = spec[model_key]["batch_size"]
@@ -56,8 +69,10 @@
             sma_columns
     )
 
-time_varying_known_categoricals = spec[model_key]["time_varying_known_categoricals"]
+else:
+    time_varying_known_reals = spec[model_key]["time_varying_known_reals"]
 
+time_varying_known_categoricals = spec[model_key]["time_varying_known_categoricals"]
 max_prediction_length = spec[model_key]["max_prediction_length"]
 max_encoder_length = spec[model_key]["max_encoder_length"]
 sample = spec[model_key]["sample"]
@@ -78,7 +93,7 @@
         train_data,
         time_idx="time_idx",
         target="value",
-        group_ids=["dataset", "id"],
+        group_ids=["id"],
         min_encoder_length=max_encoder_length // 2,
         max_encoder_length=max_encoder_length,
         min_prediction_length=1,
@@ -89,7 +104,7 @@
         time_varying_unknown_categoricals=[],
         time_varying_unknown_reals=["value"],
         target_normalizer=GroupNormalizer(
-            groups=["dataset", "id"], transformation="softplus"
+            groups=["id"], transformation="softplus"
         ),
         add_relative_time_idx=True,
         add_target_scales=True,
@@ -97,6 +112,8 @@
 
     )
 
+    pl.seed_everything(42)
+
     # create validation set (predict=True) which means to predict the
     # last max_prediction_length points in time for each series
     validation = TimeSeriesDataSet.from_dataset(
@@ -117,15 +134,6 @@
         num_workers=0
     )
 
-    # calculate baseline mean absolute error, i.e. predict next value as the
-    # last available value from the history
-    actuals = torch.cat([y for x, (y, weight) in iter(val_dataloader)])
-    baseline_predictions = Baseline().predict(val_dataloader)
-    (actuals - baseline_predictions).abs().mean().item()
-
-    # configure network and trainer
-    pl.seed_everything(42)
-
     # configure network and trainer
     early_stop_callback = EarlyStopping(
         monitor="val_loss",
@@ -134,6 +142,7 @@
         verbose=False,
         mode="min"
     )
+
     lr_logger = LearningRateMonitor()
     logger = TensorBoardLogger("lightning_logs")
 
@@ -155,7 +164,7 @@
         attention_head_size=1,
         dropout=DROPOUT,
         hidden_continuous_size=HIDDEN_CONTINUOUS_SIZE,
-        output_size=7,  # 7 quantiles by default
+        output_size=7,
         loss=QuantileLoss(),
         log_interval=10,
         reduce_on_plateau_patience=4,
@@ -168,4 +177,4 @@
         val_dataloaders=val_dataloader,
     )
 
-    torch.save(tft.state_dict(), MODEL_PATH)
\ No newline at end of file
+    torch.save(tft.state_dict(), MODEL_PATH)
diff --git a/hyperparameter_tuning.py b/tuning_tft.py
similarity index 100%
rename from hyperparameter_tuning.py
rename to tuning_tft.py