# @package _global_

# All parameters below will be merged with parameters from default configurations set above 
# this allows you to overwrite only specified parameters

# Options for training:
# - /data: [pdb, pdb_sp, pdb_sp_afesm, pdb_sp_afesme]
# - /model/architecture: 
#         [foldingdit_100M, foldingdit_360M, foldingdit_700M, 
#          foldingdit_1.1B, foldingdit_1.6B, foldingdit_3B]
# - /trainer: [default, fsdp]

defaults:
  - override /data: pdb
  - override /model: simplefold
  - override /model/architecture: foldingdit_100M
  - override /model/sampler: euler_maruyama
  - override /model/processor: protein_processor
  - override /callbacks: default
  - override /trainer: fsdp # FSDP training
  - override /logger: tensorboard

# load_ckpt_path: [PATH_TO_YOUR_CKPT] # e.g. uncomment to load a checkpoint
seed: 12345

trainer:
  max_steps: 300000
  val_check_interval: 10000 # this controls BOTH checkpoint steps and validation
  accumulate_grad_batches: 1 # always 1 for FSDP

data:
  batch_size: 8 # should be the same as the number of GPUs
  num_workers: 16

model:
  _target_: models.simplefold.simplefold.SimpleFold
  ema_decay: 0.999
  clip_grad_norm_val: 2.0
  esm_model: "esm2_3B"

  scheduler:
    _target_: onescience.utils.simplefold.lr_scheduler.LinearWarmup
    _partial_: true
    min_lr: 1e-6
    max_lr: ${model.optimizer.lr}
    warmup_steps: 5000

  architecture:
    esm_model: ${model.esm_model}

  processor:
    scale: 16.0
    ref_scale: 5.0
    multiplicity: 16 # number of copies per GPU