Update moo.py

moo.py

@@ -1,11 +1,9 @@
-import yaml
-from tqdm import tqdm
 import torch
-from torch import nn
 from transformers import AutoTokenizer
+from pathlib import Path
+import inspect
 
 from models.peptide_classifiers import *
-
 from utils.parsing import parse_guidance_args
 args = parse_guidance_args()
 
@@ -55,6 +53,19 @@ if 'Motif' in args.objectives or 'Specificity' in args.objectives:
     motif_model = MotifModel(bindevaluator, target_sequence, motifs, penalty=args.motif_penalty)
     score_models.append(motif_model)
 
+objective_line = 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:
+        with open(args.output_file, 'w') as f:
+            f.write(objective_line)
+else:
+    with open(args.output_file, 'w') as f:
+            f.write(objective_line)
+
 for i in range(args.n_batches):
     if source_distribution == "uniform":
         x_init = torch.randint(low=4, high=vocab_size, size=(n_samples, length), device=device)   # CHANGE!
@@ -87,6 +98,11 @@ 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
+
         if isinstance(candidate_scores, tuple):
             for score in candidate_scores:
                 scores.append(score.item())