configs/wt_pep.yaml

seed: 42
base_dir: /scratch/pranamlab/sgoel/MadSBM

training:
  mode: test  # train / test / resume_from_ckpt
  n_unfrozen: 3
  n_epochs: 50
  log_every_n_steps: 50
  num_sanity_val_steps: 2
  val_check_interval: 
  enable_progress_bar: true
  grad_clip_val: 10.0
  accumulate_grad_batches: 16  # to workaround dynamic batching
  devices: 1 # number of GPUs

model:
  ablate: false
  evoflow_model: fredzzp/EvoFlow-150M-afdbseq
  esm_model: facebook/esm2_t33_650M_UR50D
  n_layers: 2 #8
  n_heads: 16 #8
  hidden_dim: 1280
  attn_drop: 0.0
  resid_drop: 0.0
  mlp_ratio: 4.0
  beta1: 1e-6
  beta2: 1e-6


time_embed:
  time_dim: 512
  fourier_dim: 64
  fourier_scale: 30.0
  time_schedule: uniform  # linear / exponential / uniform
  anneal_frac: 0.75
  min_time: 1e-6
  n_timesteps: 500

data:
  batch_size: 1
  #max_seq_len: 500
  train: /scratch/pranamlab/tong/data/peptide/tokenized_peptide_batched/train
  test: /scratch/pranamlab/tong/data/peptide/tokenized_peptide_batched/test
  val: /scratch/pranamlab/tong/data/peptide/tokenized_peptide_batched/val


optim:
  type: adamw
  scheduler: cosine
  lr: 1e-4
  lr_end: 1e-5
  warmup_init_lr: 1e-6
  warmup_epochs: 2
  weight_decay: 0.01
  beta1: 0.9
  beta2: 0.999
  power: 1


wandb:
  project: MadSBM_PEPTIDE
  group: programmablebio
  name: peptide_og-madsbm_esm_no-gclip_lr=1e-4_n-layers=2_n-heads=16_trainable-lm-head_logits-sum-SM_gate-esm
  #name: peptide_og-madsbm_esm_no-gclip_lr=1e-4_n-layers=2_n-heads=16_trainable-lm-head_logits-sum-SM_ABLATE-gate-esm
  id: ${.name}_${seed}


checkpointing:
  save_every_n_epochs: 1
  save_dir: ${base_dir}/checkpoints/wt_pep/${wandb.name}
  resume_ckpt_path: ${checkpointing.save_dir}/last.ckpt
  best_ckpt_path: ${checkpointing.save_dir}/best-model_epoch=41_step=106890.ckpt


sampling:
  model_type: madsbm  # madsbm / diffusion / dfm
  n_steps: 32
  top_p: 0.9
  rate_scale: 0.01
  jump_scale: 0.05
  tau: 0.5
  M: 16
  beta: 2.0