model:
  name: "ELECTRAProteinModel"
  max_residues: 512
  embed_dim: 2560
  encoder_depth: 33
  encoder_heads: 40
  decoder_dim: 2560
  decoder_depth: 3
  decoder_heads: 40
  generator_depth: 8
  generator_decoder_depth: 2
  disc_loss_weight: 50.0         # corrective CE same scale as generator CE
  generator_lr_multiplier: 1.0
  generator_temperature: 2.0
  relax_temperature_scaling: null   # null = 1/sqrt(head_dim), relaxes attention temperature
  use_qk_norm: false

# Training configuration
training:
  learning_rate: 1.0e-3
  weight_decay: 0.01
  adam_beta1: 0.9
  adam_beta2: 0.95
  warmup_steps: 5000
  batch_size: 32
  num_workers: 8
  max_steps: 500000
  total_training_steps: 500000
  accumulate_grad_batches: 1
  val_check_interval: 100
  test_interval: 100
  gradient_clip_val: 1.0

# Data configuration
data:
  name: "ELECTRADataModule"
  lmdb_dir: "/path/to/triprorep-pretrain"
  struct_format: "fullatom"
  max_length: 512
  bb_vocab_size: 512
  fa_vocab_size: 512
  mask_prob: 0.6
  random_cropping: true
  pack_sequences: false

# Lightning configuration
lightning:
  gpus: 8
  nodes: 4
  precision: "bf16-mixed"
  compile: true
  checkpoint_dir: "./checkpoints"
  log_dir: "./logs"
  wandb_dir: "./wandb"
  wandb_project: "k-fold-structure"
  wandb_name: "triprorep_3B"
  save_every_n_steps: 10000