a2d2_pep/config_pep.yaml

trainer: "any-order-flow"
dataset: "peptides"

model:
  hidden_size: 768
  n_heads: 12
  cond_dim: 128
  dropout: 0.05
  n_blocks: 12

interpolant:
  type: "any-order"
  tokens: null # filled in automatically
  pad_token: null # filled in automatically
  mask_token: null # filled in automatically
  max_length: 1024
  insert_schedule:
    type: "linear"
  unmask_schedule:
    type: "linear"

training:
  only_embed_insert: true
  batch_size: 1024
  per_gpu_batch_size: 64 # Gradient accumulation happens automatically
  cpus: 4
  learning_rate: 3e-4
  nodes: 1
  devices: 4
  max_steps: 1000000
  weight_decay: 0.03
  # Path to the preprocessed (arrow) pretraining dataset; see README for the download link.
  # Relative paths resolve against a2d2_pep/. Defaults to a2d2_pep/data/11M_peptide_smiles.
  data_path: "data/11M_peptide_smiles"
  checkpoint_dir: "checkpoints/peptides"
  save_top_k: 1
  save_every_n_epochs: 1
  loss_fn:
    unmask: "elbo"
    insert: "expectation"
  reset_lr: false
  warmup_steps: 2000
  ema_decay: 0.9999
  filter_max_length: false

wandb:
  entity: null  # set to your W&B entity, or leave null to use the default
  project: "a2d2-pep"
  name: "a2d2-pep"
  path: "./wandb"