unimol_args:
  data: ./data/geom
  dict_name: dict.txt
  user_dir: ./models/unimol/unimol
  train_subset: train
  valid_subset: valid
  data_path: ./data/geom
  save_dir: ./unimol_pretrain_save

  n_gpu: 2
  master_port: 10086
  num_workers: 8
  ddp_backend: c10d

  task: unimol
  loss: unimol
  arch: unimol_base

  optimizer: adam
  adam_betas: [0.9, 0.99]
  adam_eps: 1e-6
  clip_norm: 1.0
  weight_decay: 1e-4

  lr_scheduler: polynomial_decay
  lr: 1e-4
  warmup_updates: 10000
  total_num_update: 1000000

  update_freq: 1
  seed: 1

  fp16: true
  fp16_init_scale: 4
  fp16_scale_window: 256

  max_update: 1000000
  log_interval: 10
  log_format: simple

  save_interval_updates: 10000
  validate_interval_updates: 10000
  keep_interval_updates: 10
  no_epoch_checkpoints: true

  masked_token_loss: 1
  masked_coord_loss: 5
  masked_dist_loss: 10
  x_norm_loss: 0.01
  delta_pair_repr_norm_loss: 0.01
  mask_prob: 0.15
  only_polar: -1  # Use all H

  noise_type: uniform
  noise: 1.0
  batch_size: 64

  mode: infer