Update moo.py

moo.py

@@ -33,17 +33,27 @@ 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)
 
-bindevaluator = load_bindevaluator('./classifier_ckpt/finetuned_BindEvaluator.ckpt', device)
-motif_model = MotifModel(bindevaluator, target_sequence, motifs, penalty=args.motif_penalty)
-
-affinity_predictor = load_affinity_predictor('./classifier_ckpt/binding_affinity_unpooled.pt', device)
-affinity_model = AffinityModel(affinity_predictor, target_sequence)
-hemolysis_model = HemolysisModel(device=device)
-nonfouling_model = NonfoulingModel(device=device)
-solubility_model = SolubilityModelNew(device=device)
-halflife_model = HalfLifeModel(device=device)
-
-score_models = [hemolysis_model, nonfouling_model, halflife_model, affinity_model, motif_model]
+score_models = []
+if 'Hemolysis' in args.objectives:
+    hemolysis_model = HemolysisModel(device=device)
+    score_models.append(hemolysis_model)
+if 'Non-Fouling' in args.objectives:
+    nonfouling_model = NonfoulingModel(device=device)
+    score_models.append(nonfouling_model)
+if 'Solubility' in args.objectives:
+    solubility_model = SolubilityModelNew(device=device)
+    score_models.append(solubility_model)
+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('/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('/scratch/pranamlab/tong/checkpoints/BindEvaluator/model_path/finetuned_BindEvaluator.ckpt', device)
+    motif_model = MotifModel(bindevaluator, target_sequence, motifs, penalty=args.motif_penalty)
+    score_models.append(motif_model)
 
 for i in range(args.n_batches):
     if source_distribution == "uniform":