dvc.yaml

stages:
  load_dataset:
    cmd: python ./src/data/load_dataset.py
    deps:
    - ./src/data/load_dataset.py
    params:
    - load_dataset.bucket
    - load_dataset.raw_data
    - load_dataset.filename
    outs:
    - .${load_dataset.raw_data}/${load_dataset.filename}

  make_dataset:
    cmd: python ./src/data/make_dataset.py
    deps:
    - ./src/data/make_dataset.py
    - .${load_dataset.raw_data}/${load_dataset.filename}
    params:
    - make_dataset.input_path
    - make_dataset.holdout_split
    - make_dataset.unit_test_split
    - make_dataset.seed
    - make_dataset.train_test_path
    - make_dataset.unit_test_path
    outs:
    - .${make_dataset.train_test_path}/train.csv
    - .${make_dataset.train_test_path}/test.csv
    - .${make_dataset.unit_test_path}/unit_test.csv

  feature_engineering:
    cmd: python ./src/features/feature_engineering.py
    deps:
    - ./src/features/feature_engineering.py
    - .${make_dataset.train_test_path}/train.csv
    - .${make_dataset.train_test_path}/test.csv
    params:
    - feature_engineering.input_path__train
    - feature_engineering.input_path__test
    - feature_engineering.export_path
    - feature_engineering.labels_
    - feature_engineering.filename_train
    - feature_engineering.filename_test
    outs:
    - .${feature_engineering.export_path}/${feature_engineering.filename_train}.csv
    - .${feature_engineering.export_path}/${feature_engineering.filename_test}.csv

  train_model:
    cmd: python ./src/models/train_model.py
    deps: 
    - ./src/models/train_model.py
    - .${feature_engineering.export_path}/${feature_engineering.filename_train}.csv
    - .${feature_engineering.export_path}/${feature_engineering.filename_test}.csv
    params:
    - base.target
    - feature_engineering.export_path
    - mlflow.repo_name
    - mlflow.repo_owner
    - train_model.model_to_train
    - train_model.model_name
    - train_model.decision_tree
    - train_model.random_forest
    - train_model.gradient_boost
    - train_model.xgb
    outs:
    - ./models/${train_model.model_name}.joblib

  tune_model:
    cmd: python ./src/models/tune_model.py
    deps:
    - ./src/models/tune_model.py
    - .${feature_engineering.export_path}/${feature_engineering.filename_train}.csv
    - .${feature_engineering.export_path}/${feature_engineering.filename_test}.csv
    params:
    - base.target
    - feature_engineering.export_path
    - mlflow.repo_name
    - mlflow.repo_owner
    - tune_model.model_name
    - tune_model.n_trials
    outs:
    - ./models/${tune_model.model_name}.joblib