dvc.yaml

params:
  - params.yaml
#   - params.yaml
  # - conf
  # - conf/config
  # - conf/config.yaml
  # - config.yaml
  # - task: aq
  # - split: scaffold
  # - model: mmb
  # - head: hier
    #- config.yaml
    #- cfg/ds
  

    #- cfg/data_cfg.json:property,split,data_seed
    #- cfg/model_cfg.json:

#   - split: random,accurate,scaffold
#   - model: mmb-lin,mmb-hier,mmb-ft-hier,mmb-ft-lin,mmb-avg,ecfp-hier,ecfp-lin

stages:
  split_data:
    # matrix:
    #   split: ${data.split}
    cmd: 
      # - echo "${item.task} ${item.split}"
      - pwd
      - echo "SPLIT DATA ${task} ${split}"
      - mkdir -p data/${task.task}/${split.split}/
      - python3 scripts/split_data.py #${task} ${split}
    params:
      - task
      - split
    deps:
      # - data/AqueousSolu.csv
      - src/dataloader.py
      - scripts/split_data.py
    outs:
      - data/${task.task}/${split.split}/:
          persist: true

  train_model:
    cmd: 
      - mkdir -p out/${task.task}/${split.split}/${model.model}-${head.head}/{viz,model}
      - echo "TRAIN ${model} ${head}"
      # - python3 scripts/train_model.py #${task} ${split} ${model} ${head}
      - python3 scripts/train_${head.fit}.py #${task} ${split} ${model} ${head}
      
    params:
      - task
      - split
      - model
      - head
    deps:
      - data/${task.task}/${split.split}/
      - src/model.py
      - scripts/train_model.py
      - scripts/train_sklearn.py
    metrics:
      - out/${task.task}/${split.split}/${model.model}-${head.head}/metrics.json:
          persist: true
    outs:
      - out/${task.task}/${split.split}/${model.model}-${head.head}/model/:
          persist: true
      - out/${task.task}/${split.split}/${model.model}-${head.head}/best.pt:
          persist: true

  predict_model:
    cmd: 
      - mkdir -p out/${task.task}/${split.split}/${model.model}-${head.head}/viz
      - python3 scripts/predict_model.py
      - python3 scripts/plot_datasplit.py
      # - cp -r out/${task.task}/${split.split}/${model.model}-${head.head}/
       # /workspace/final/${task.task}/${split.split}/
    deps:
      - src/model.py
      - scripts/predict_model.py
      - scripts/plot_datasplit.py
      - data/${task.task}/${split.split}/
      - out/${task.task}/${split.split}/${model.model}-${head.head}/best.pt
    params:
      - task
      - split
      - model
      - head
    outs:
      - out/${task.task}/${split.split}/${model.model}-${head.head}/parity_plot.png:
          persist: true
      - out/${task.task}/${split.split}/${model.model}-${head.head}/latent_viz.png:
          persist: true
      - out/${task.task}/${split.split}/${model.model}-${head.head}/predictions.csv:
          persist: true

  explain_model:
    cmd:
      # - rm out/${task.task}/${split.split}/${model.model}-${head.head}/viz/*
      - mkdir -p out/${task.task}/${split.split}/${model.model}-${head.head}/viz
      - python3 scripts/explain_${xai.xai}.py #${task} ${split} ${model} ${head} 
    deps:
      - scripts/explain_${xai.xai}.py
      - data/${task.task}/${split.split}/test.pkl
      - out/${task.task}/${split.split}/${model.model}-${head.head}/best.pt
      - src/maskedhead.py
      - src/explainer.py
    params:
      - task
      - split
      - model
      - head
      - xai
    outs:
      - out/${task.task}/${split.split}/${model.model}-${head.head}/viz:
          persist: true
      - out/${task.task}/${split.split}/${model.model}-${head.head}/attributions.csv:
          persist: true

  all:
    cmd: 
      - pwd
      - mkdir -p /workspace/final/${task.task}/${split.split}/${model.model}-${head.head}/{model,viz}
      - cp -r out/${task.task}/${split.split}/${model.model}-${head.head}/
       /workspace/final/${task.task}/${split.split}/
      - cp -r dvc.lock /workspace/final/${task.task}/${split.split}/${model.model}-${head.head}/dvc.lock
      - cp -r params.yaml /workspace/final/${task.task}/${split.split}/${model.model}-${head.head}/params.yaml
      - python3 scripts/plot_similarity.py
      - python3 scripts/plot_models.py
      - echo "Finished full ${task.task} ${split.split} pipeline with ${model.model}-${head.head}"
    deps:
      - scripts/plot_models.py
      - scripts/plot_similarity.py
      - out/${task.task}/${split.split}/
      - out/${task.task}/${split.split}/${model.model}-${head.head}/model
      - out/${task.task}/${split.split}/${model.model}-${head.head}/best.pt
      - out/${task.task}/${split.split}/${model.model}-${head.head}/parity_plot.png
      - out/${task.task}/${split.split}/${model.model}-${head.head}/latent_viz.png
      - out/${task.task}/${split.split}/${model.model}-${head.head}/predictions.csv
      - out/${task.task}/${split.split}/${model.model}-${head.head}/attributions.csv
      - out/${task.task}/${split.split}/${model.model}-${head.head}/viz
      # - final/${task.task}/${split.split}/${model.model}-${head.head}/model_comparison_mae.png
      # - final/${task.task}/${split.split}/${model.model}-${head.head}/model_comparison_rmse.png