Merge pull request #34 from YerevaNN/mol_opt

refactor mol_opt

chemlactica/mol_opt/chemlactica_125m_hparams.yaml

@@ -1,15 +1,12 @@
-# checkpoint_path: /nfs/dgx/raid/chem/checkpoints/facebook/galactica-125m/9954e52e400b43d18d3a40f6/checkpoint-20480
-# checkpoint_path: /nfs/dgx/raid/chem/checkpoints/facebook/galactica-125m/9954e52e400b43d18d3a40f6/checkpoint-12288
-checkpoint_path: /home/admin/checkpoints/facebook/galactica-125m/1f289ff103034364bd27e1c3/checkpoint-18000
-tokenizer_path: /home/admin/tigran/ChemLactica/chemlactica/tokenizer/ChemLacticaTokenizer66
-pool_size: 50
+checkpoint_path: /nfs/dgx/raid/chem/checkpoints/facebook/galactica-125m/1f289ff103034364bd27e1c3/checkpoint-18000
+tokenizer_path: ChemLactica/chemlactica/tokenizer/ChemLacticaTokenizer66
+pool_size: 10
 validation_perc: 0.2
 num_mols: 0
-num_similars: 1
+num_similars: 5
 num_gens_per_iter: 200
 device: cuda:0
-sim_range: [0.8, 0.9]
-# qed_range: [0.5, 0.9]
+sim_range: [0.4, 0.9]
 num_processes: 8
 generation_batch_size: 200
 eos_token: "</s>"
@@ -21,18 +18,19 @@ generation_config:
   do_sample: true
   eos_token_id: 20
 
-strategy: [default]
+strategy: [rej-sample-v2]
 
 rej_sample_config:
   train_tol_level: 3
-  checkpoints_dir: ./
-  max_learning_rate: 0.00001
+  checkpoints_dir: checkpoints
+  max_learning_rate: 0.0001
+  lr_end: 0
   train_batch_size: 2
   gradient_accumulation_steps: 8
   weight_decay: 0.1
   adam_beta1: 0.9
   adam_beta2: 0.999
-  warmup_steps: 0
+  warmup_steps: 10
   global_gradient_norm: 1.0
   dataloader_num_workers: 1
   max_seq_length: 2048