Add large file using Git LFS

configs/config.yaml

@@ -0,0 +1,171 @@
+defaults:
+  - path
+
+noise:
+  type: loglinear
+  sigma_min: 1e-4
+  sigma_max: 20
+  state_dependent: True
+
+mode: ppl_eval  # train / ppl_eval / sample_eval
+diffusion: absorbing_state
+vocab: old_smiles # old_smiles / new_smiles / selfies / helm
+backbone: roformer  # peptideclm / helmgpt / dit / roformer / finetune_roformer
+parameterization: subs  # subs
+time_conditioning: False
+T: 0  # 0 (continuous time) / 1000 
+subs_masking: False
+
+seed: 42
+
+mcts: 
+  num_children: 50
+  num_objectives: 5
+  topk: 100
+  mask_token: 4
+  num_iter: 128
+  sampling: 0 # 0 is gumbel sampling / > 0 samples children from top k probs
+  invalid_penalty: 0.5
+  sample_prob: 1.0
+  perm: True
+  dual: False
+  single: False
+  time_dependent: True
+
+lr_scheduler:
+  _target_: transformers.get_constant_schedule_with_warmup
+  num_warmup_steps: 2500
+
+data:
+  train: ${paths.data}/finetune2/30K-train.csv
+  valid: ${paths.data}/finetune2/30K-val.csv
+  batchinohup ng: wrapping # padding / wrapping
+
+loader:
+  global_batch_size: 64
+  eval_global_batch_size: ${.global_batch_size}
+  # Note: batch_size and eval_batch_size are **per machine**
+  batch_size: ${div_up:${.global_batch_size}, ${eval:${trainer.devices} * ${trainer.num_nodes}}}
+  eval_batch_size: ${div_up:${.eval_global_batch_size}, ${eval:${trainer.devices} * ${trainer.num_nodes}}}
+  num_workers: ${eval:"len(__import__('os').sched_getaffinity(0))"}
+  pin_memory: True
+
+sampling:
+  predictor: ddpm_cache  # analytic, ddpm, ddpm_cache
+  num_sequences: 100
+  sampling_eps: 1e-3
+  steps: 128
+  seq_length: 100
+  noise_removal: True
+  num_sample_batches: 2  # Total samples: `num_gpus` * `loader.eval_batch_size` * num_sample_batches
+  num_sample_log: 2
+  stride_length: 1
+  num_strides: 1
+
+training:
+  antithetic_sampling: True
+  sampling_eps: 1e-3
+  focus_mask: False
+  #dynamic_batching: True
+  accumulator: False
+
+eval:
+  checkpoint_path: ${paths.checkpoints}/11M-old-tokenizer/epoch=10-step=156276.ckpt
+  disable_ema: False
+  compute_generative_perplexity: False
+  perplexity_batch_size: 8
+  compute_perplexity_on_sanity: False
+  gen_ppl_eval_model_name_or_path: gpt2-large  # gpt2-large, meta-llama/Llama-2-7b-hf
+  generate_samples: True
+  generation_model: ${paths.checkpoints}/11M-old-tokenizer/
+  
+optim:
+  weight_decay: 0.075
+  lr: 3e-4
+  beta1: 0.9
+  beta2: 0.999
+  eps: 1e-8
+
+pepclm:
+  hidden_size: 768
+  cond_dim: 256
+  n_heads: 20
+  n_blocks: 4
+  dropout: 0.5
+  length: 512
+  #scale_by_sigma: True
+
+model:
+  type: ddit
+  hidden_size: 768
+  cond_dim: 128
+  length: 512
+  n_blocks: 12
+  n_heads: 12
+  scale_by_sigma: True
+  dropout: 0.1
+
+roformer:
+  hidden_size: 768
+  n_layers: 8
+  n_heads: 8
+  max_position_embeddings: 1035
+
+helmgpt:
+  hidden_size: 256
+  embd_pdrop: 0.1
+  resid_pdrop: 0.1
+  attn_pdrop: 0.1
+  ff_dropout: 0.
+  block_size: 140
+  n_layer: 8
+  n_heads: 8
+
+
+trainer:
+  _target_: lightning.Trainer
+  accelerator: cuda
+  num_nodes: 1
+  devices: ${device_count:}
+  accumulate_grad_batches: ${div_up:${loader.global_batch_size}, ${eval:${trainer.devices} * ${loader.batch_size} * ${trainer.num_nodes}}}
+  gradient_clip_val: 1.0
+  precision: 64-true
+  num_sanity_val_steps: 2
+  max_epochs: 100
+  max_steps: 1_000_000
+  log_every_n_steps: 10
+  limit_train_batches: 1.0   # train on full dataset, can be used to toggle quick run
+  limit_val_batches: 1.0     # validate on full dataset, can be used to toggle quick run
+  #val_check_interval: 40 #954
+  check_val_every_n_epoch: 1
+
+
+wandb:
+  project: ${env_or:WANDB_PROJECT,peptune}
+  notes: null
+  group: null
+  job_type: null
+  name: ${env_or:WANDB_RUN_NAME,local}
+  id: ${.name}
+
+hydra:
+  run:
+    dir: ./${now:%Y.%m.%d}/
+  job:
+    chdir: True
+
+checkpointing:
+  # Use custom `save_dir` if, e.g., saving to S3 bucket, otherwise leave this parameter as is
+  save_dir: ${paths.outputs}
+  # Note: `checkpoints` path should correspond to `checkpoint_every_n_steps.dirpath`
+  resume_from_ckpt: True
+  resume_ckpt_path: ${paths.checkpoints}/last.ckpt
+
+callbacks:
+  model_checkpoint:
+    _target_: pytorch_lightning.callbacks.ModelCheckpoint
+    every_n_epochs: 1
+    monitor: "val/nll"
+    save_top_k: 10
+    mode: "min"
+    dirpath: ${paths.checkpoints}/11M-old-tokenizer

configs/path.yaml

@@ -0,0 +1,7 @@
+paths:
+  base: ${cwd:}             # repo root
+  data: ${paths.base}/data
+  checkpoints: ${paths.base}/checkpoints
+  tokenizers: ${paths.base}/tokenizers
+  outputs: ${paths.base}/outputs
+

main.py

@@ -16,7 +16,6 @@ import torch
 import sys
 import torch.distributed as dist
 from torch.nn.parallel import DistributedDataParallel as DDP
-sys.path.append("/home/st512/peptune/scripts/peptide-mdlm-mcts")
 
 import dataset as dataloader
 import dataloading_for_dynamic_batching as dynamic_dataloader
@@ -30,24 +29,25 @@ from tokenizer.my_tokenizers import SMILES_SPE_Tokenizer
 from helm_tokenizer.helm_tokenizer import HelmTokenizer
 
 
-#wandb.login(key="5a7613c531cb58f9802f3f8e2f73bc4997b917ab")
-
 omegaconf.OmegaConf.register_new_resolver('cwd', os.getcwd)
 omegaconf.OmegaConf.register_new_resolver('device_count', torch.cuda.device_count)
 omegaconf.OmegaConf.register_new_resolver('eval', eval)
 omegaconf.OmegaConf.register_new_resolver('div_up', lambda x, y: (x + y - 1) // y)
+omegaconf.OmegaConf.register_new_resolver("env_or", lambda k, d: os.getenv(k, d))
 
 def _load_from_checkpoint(config, tokenizer):
-	if 'hf' in config.backbone:
-		return Diffusion(
-			config, tokenizer=tokenizer).to('cuda')
-	else:
-		model = Diffusion.load_from_checkpoint(
-			config.eval.checkpoint_path,
-			tokenizer=tokenizer,
-			config=config)
-
-	return model
+    """Create Diffusion model; load weights if checkpoint_path is set."""
+    if "hf" in str(config.get("backbone", "")):
+        return Diffusion(config, tokenizer=tokenizer).to("cuda")
+
+    ckpt_path = config.eval.checkpoint_path
+    model = Diffusion.load_from_checkpoint(
+        ckpt_path,
+        tokenizer=tokenizer,
+        config=config,
+        map_location="cuda" if torch.cuda.is_available() else "cpu",
+    )
+    return model
 
 @L.pytorch.utilities.rank_zero_only
 def print_config(
@@ -197,36 +197,37 @@ def _train(config, logger, tokenizer, data_module):
 	
 	model = Diffusion(config, tokenizer=tokenizer)
 	
-	if config.backbone == 'finetune_roformer':
-		checkpoint = torch.load('/home/st512/peptune/scripts/peptide-mdlm-mcts/checkpoints/11M-old-tokenizer/epoch=1-step=24080.ckpt')
-		model.load_state_dict(checkpoint['state_dict'])
+	if config.backbone == 'finetune_roformer' and config.eval.checkpoint_path:
+		checkpoint = torch.load(config.eval.checkpoint_path, map_location="cpu")
+        state = checkpoint.get("state_dict", checkpoint)
+		model.load_state_dict(state, strict=False)
 	
 	trainer.fit(model, datamodule=data_module, ckpt_path=ckpt_path)
 
   
-@hydra.main(version_base=None, config_path='/home/st512/peptune/scripts/peptide-mdlm-mcts', config_name='config')
+@hydra.main(version_base=None, config_path='configs', config_name='config')
 def main(config):
 	"""
  		Main entry point for training
    """   
-	wandb.init(project="peptune")
 	L.seed_everything(config.seed)
  
 	# print_config(config, resolve=True, save_cfg=True)
 
 	logger = utils.get_logger(__name__)
 	# load PeptideCLM tokenizer
-	if config.vocab == 'new_smiles':
+	tok_dir = config.paths.tokenizers
+    if config.vocab == 'new_smiles':
 		tokenizer = APETokenizer()
-		tokenizer.load_vocabulary('/home/st512/peptune/scripts/peptide-mdlm-mcts/new_tokenizer/peptide_smiles_600_vocab.json')
+		tokenizer.load_vocabulary(f'{tok_dir}/peptide_smiles_600_vocab.json')
 	elif config.vocab == 'old_smiles':
-		tokenizer = SMILES_SPE_Tokenizer('/home/st512/peptune/scripts/peptide-mdlm-mcts/tokenizer/new_vocab.txt', 
-                                   '/home/st512/peptune/scripts/peptide-mdlm-mcts/tokenizer/new_splits.txt')
+		tokenizer = SMILES_SPE_Tokenizer(f'{tok_dir}/new_vocab.txt', 
+                                   f'{tok_dir}/new_splits.txt')
 	elif config.vocab == 'selfies':
 		tokenizer = APETokenizer()
-		tokenizer.load_vocabulary('/home/st512/peptune/scripts/peptide-mdlm-mcts/new_tokenizer/peptide_selfies_600_vocab.json')
+		tokenizer.load_vocabulary(f'{tok_dir}/peptide_selfies_600_vocab.json')
 	elif config.vocab == 'helm':
-		tokenizer = HelmTokenizer('/home/st512/peptune/scripts/peptide-mdlm-mcts/helm_tokenizer/monomer_vocab.txt')
+		tokenizer = HelmTokenizer(f'{tok_dir}/monomer_vocab.txt')
 
 	if config.backbone == 'finetune_roformer':
 		train_dataset = load_dataset('csv', data_files=config.data.train)
@@ -236,7 +237,7 @@ def main(config):
 		val_dataset = val_dataset['train']#.select(lst)
 		data_module = dataloader.CustomDataModule(train_dataset, val_dataset, None, tokenizer, batch_size=config.loader.global_batch_size)
 	else:
-		data_module = dynamic_dataloader.CustomDataModule('/home/st512/peptune/scripts/peptide-mdlm-mcts/data/smiles/11M_smiles_old_tokenizer_no_limit', tokenizer)
+		data_module = dynamic_dataloader.CustomDataModule(f'{config.paths.data}/smiles/11M_smiles_old_tokenizer_no_limit', tokenizer)
 	
 	if config.mode == 'sample_eval':
 		generate_samples(config, logger, tokenizer)
@@ -247,4 +248,4 @@ def main(config):
 
 
 if __name__ == '__main__':
-	main()
+	main()

peptune-pretrained.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b259f022c21121f5c755fed61230d6fdf2626ee4ab8a23df479b3cf553fd4aef
+size 1386966244