Update moppit.py

moppit.py

@@ -22,7 +22,7 @@ if args.motifs:
     print(motifs)
 
 tokenizer = AutoTokenizer.from_pretrained("facebook/esm2_t33_650M_UR50D")
-target_sequence = tokenizer(target, return_tensors='pt')['input_ids'].to(device)
+target_sequence = tokenizer(target, return_tensors='pt').to(device)
 
 # Load Models
 solver = load_solver('./ckpt/peptide/cnn_epoch200_lr0.0001_embed512_hidden256_loss3.1051.ckpt', vocab_size, device)
@@ -35,14 +35,14 @@ 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)
+    solubility_model = SolubilityModel(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('./classifier_ckpt/binding_affinity_unpooled.pt', device)
-    affinity_model = AffinityModel(affinity_predictor, target_sequence)
+    affinity_predictor = load_affinity_predictor('../classifier_ckpt/wt_affinity.pt', device)
+    affinity_model = AffinityModel(affinity_predictor, target_sequence, device)
     score_models.append(affinity_model)
 if 'Motif' in args.objectives or 'Specificity' in args.objectives:
     bindevaluator = load_bindevaluator('./classifier_ckpt/finetuned_BindEvaluator.ckpt', device)
@@ -50,7 +50,7 @@ if 'Motif' in args.objectives or 'Specificity' in args.objectives:
         motif_penalty = True
     else: 
         motif_penalty = False
-    motif_model = MotifModel(bindevaluator, target_sequence, motifs, penalty=motif_penalty)
+    motif_model = MotifModel(bindevaluator, target_sequence['input_ids'], motifs, penalty=motif_penalty)
     score_models.append(motif_model)
 
 objective_line = "Binder," + str(args.objectives)[1:-1].replace(' ', '').replace("'", "") + '\n'