Update moppit.py

moppit.py

@@ -7,10 +7,6 @@ from models.peptide_classifiers import *
 from utils.parsing import parse_guidance_args
 args = parse_guidance_args()
 
-import pdb
-import random
-import inspect
-import csv
 
 # MOO hyper-parameters
 step_size = 1 / 100
@@ -29,7 +25,7 @@ tokenizer = AutoTokenizer.from_pretrained("facebook/esm2_t33_650M_UR50D")
 target_sequence = tokenizer(target, return_tensors='pt')['input_ids'].to(device)
 
 # Load Models
-solver = load_solver('./ckpt/peptide/cnn_epoch200_lr0.0001_embed512_hidden256_loss3.1051.ckpt', vocab_size, device)
+solver = load_solver('/scratch/pranamlab/tong/checkpoints/MOG-DFM/ckpt/peptide/cnn_epoch200_lr0.0001_embed512_hidden256_loss3.1051.ckpt', vocab_size, device)
 
 score_models = []
 if 'Hemolysis' in args.objectives:
@@ -45,21 +41,25 @@ if 'Half-Life' in args.objectives:
     halflife_model = HalfLifeModel(device=device)
     score_models.append(halflife_model)
 if 'Affinity' in args.objectives:
-    affinity_predictor = load_affinity_predictor('./classifier_ckpt/binding_affinity_unpooled.pt', device)
+    affinity_predictor = load_affinity_predictor('/scratch/pranamlab/tong/checkpoints/MOG-DFM/classifier_ckpt/binding_affinity_unpooled.pt', device)
     affinity_model = AffinityModel(affinity_predictor, target_sequence)
     score_models.append(affinity_model)
 if 'Motif' in args.objectives or 'Specificity' in args.objectives:
-    bindevaluator = load_bindevaluator('./classifier_ckpt/finetuned_BindEvaluator.ckpt', device)
-    motif_model = MotifModel(bindevaluator, target_sequence, motifs, penalty=args.motif_penalty)
+    bindevaluator = load_bindevaluator('/scratch/pranamlab/tong/checkpoints/BindEvaluator/model_path/finetuned_BindEvaluator.ckpt', device)
+    if 'Specificity' in args.objectives:
+        motif_penalty = True
+    else: 
+        motif_penalty = False
+    motif_model = MotifModel(bindevaluator, target_sequence, motifs, penalty=motif_penalty)
     score_models.append(motif_model)
 
-objective_line = str(args.objectives)[1:-1].replace(' ', '').replace("'", "") + '\n'
+objective_line = "Binder," + str(args.objectives)[1:-1].replace(' ', '').replace("'", "") + '\n'
 
 if Path(args.output_file).exists():
     with open(args.output_file, 'r') as f:
         lines = f.readlines()
 
-    if lines[0] != objective_line:
+    if len(lines) == 0 or lines[0] != objective_line:
         with open(args.output_file, 'w') as f:
             f.write(objective_line)
 else:
@@ -83,7 +83,7 @@ for i in range(args.n_batches):
                                       verbose=True, 
                                       time_grid=torch.tensor([0.0, 1.0-1e-3]),
                                       score_models=score_models,
-                                      num_objectives=len(score_models) + int(args.motif_penalty),
+                                      num_objectives=len(score_models) + int(motif_penalty),
                                       weights=args.weights)
     
     samples = x_1.tolist()
@@ -98,8 +98,6 @@ for i in range(args.n_batches):
         else:
             candidate_scores = s(x_1)
 
-        if args.objectives[i] == 'Half-Life':
-            candidate_scores = 10 ** (candidate_scores * 2)
         if args.objectives[i] == 'Affinity':
             candidate_scores = 10 * candidate_scores