Update app.py

app.py

@@ -4,138 +4,54 @@ from utils import create_vocab, setup_seed
 from dataset_mlm import  get_paded_token_idx_gen, add_tokens_to_vocab
 setup_seed(4)
 
-# def CTXGen(X1,X2,X3,model_name):
-#     device = torch.device("cpu")
-#     vocab_mlm = create_vocab()
-#     vocab_mlm = add_tokens_to_vocab(vocab_mlm)
-#     save_path = model_name
-#     model = torch.load(save_path, weights_only=False, map_location=torch.device('cpu'))
-#     model = model.to(device)
-    
-#     predicted_token_probability_all = []
-#     model.eval()
-#     topk = []
-#     with torch.no_grad():
-#         new_seq = None
-#         seq = [f"{X1}|{X2}|{X3}|||"]
-#         vocab_mlm.token_to_idx["X"] = 4
-#         padded_seq, _, idx_msa, _ = get_paded_token_idx_gen(vocab_mlm, seq, new_seq)
-#         idx_msa = torch.tensor(idx_msa).unsqueeze(0).to(device)
-#         mask_positions = [i for i, token in enumerate(padded_seq) if token == "X"]
-#         #if not mask_positions:
-#             #raise ValueError("Nothing found in the sequence to predict.")
-
-#         for mask_position in mask_positions:
-#             padded_seq[mask_position] = "[MASK]"
-#             input_ids = vocab_mlm.__getitem__(padded_seq)
-#             input_ids = torch.tensor([input_ids]).to(device)
-#             logits = model(input_ids, idx_msa)
-#             mask_logits = logits[0, mask_position, :]
-#             predicted_token_probability, predicted_token_id = torch.topk((torch.softmax(mask_logits, dim=-1)), k=5)
-#             topk.append(predicted_token_id)
-#             predicted_token = vocab_mlm.idx_to_token[predicted_token_id[0].item()]
-#             predicted_token_probability_all.append(predicted_token_probability[0].item())
-#             padded_seq[mask_position] = predicted_token
-
-#         cls_pos = vocab_mlm.to_tokens(list(topk[0]))
-#         if X1 != "X":
-#             Topk = X1
-#             Subtype = X1
-#             Potency = padded_seq[2],predicted_token_probability_all[0]
-#         elif X2 != "X":
-#             Topk = cls_pos
-#             Subtype = padded_seq[1],predicted_token_probability_all[0]
-#             Potency = X2
-#         else:
-#             Topk = cls_pos
-#             Subtype = padded_seq[1],predicted_token_probability_all[0]
-#             Potency = padded_seq[2],predicted_token_probability_all[1]
-#     return Subtype, Potency, Topk
-
-def CTXGen(X1, X2, X3, model_name):
+def CTXGen(X1,X2,X3,model_name):
     device = torch.device("cpu")
     vocab_mlm = create_vocab()
     vocab_mlm = add_tokens_to_vocab(vocab_mlm)
-
-    model = torch.load(model_name, weights_only=False, map_location=device)
+    save_path = model_name
+    model = torch.load(save_path, weights_only=False, map_location=torch.device('cpu'))
     model = model.to(device)
-    model.eval()
-
-    seq = [f"{X1}|{X2}|{X3}|||"]
-    vocab_mlm.token_to_idx["X"] = 4
-
-    padded_seq, _, idx_msa, _ = get_paded_token_idx_gen(vocab_mlm, seq, None)
-    idx_msa = torch.tensor(idx_msa).unsqueeze(0).to(device)
-
-    mask_positions = [i for i, token in enumerate(padded_seq) if token == "X"]
-
+    
     predicted_token_probability_all = []
+    model.eval()
     topk = []
-
     with torch.no_grad():
-        if mask_positions:
-            for mask_position in mask_positions:
-                padded_seq[mask_position] = "[MASK]"
-                input_ids = torch.tensor(
-                    [vocab_mlm.__getitem__(padded_seq)]
-                ).to(device)
-
-                logits = model(input_ids, idx_msa)
-                mask_logits = logits[0, mask_position, :]
-                probs = torch.softmax(mask_logits, dim=-1)
-
-                prob, token_id = torch.topk(probs, k=5)
-                topk.append(token_id)
-
-                predicted_token = vocab_mlm.idx_to_token[token_id[0].item()]
-                predicted_token_probability_all.append(prob[0].item())
-
-                padded_seq[mask_position] = predicted_token
-
-            cls_pos = vocab_mlm.to_tokens(list(topk[0]))
-
-            if X1 != "X":
-                Subtype = X1
-                Potency = (padded_seq[2], predicted_token_probability_all[0])
-            elif X2 != "X":
-                Subtype = (padded_seq[1], predicted_token_probability_all[0])
-                Potency = X2
-            else:
-                Subtype = (padded_seq[1], predicted_token_probability_all[0])
-                Potency = (padded_seq[2], predicted_token_probability_all[1])
-
+        new_seq = None
+        seq = [f"{X1}|{X2}|{X3}|||"]
+        vocab_mlm.token_to_idx["X"] = 4
+        padded_seq, _, idx_msa, _ = get_paded_token_idx_gen(vocab_mlm, seq, new_seq)
+        idx_msa = torch.tensor(idx_msa).unsqueeze(0).to(device)
+        mask_positions = [i for i, token in enumerate(padded_seq) if token == "X"]
+        if not mask_positions:
+            raise ValueError("Nothing found in the sequence to predict.")
+
+        for mask_position in mask_positions:
+            padded_seq[mask_position] = "[MASK]"
+            input_ids = vocab_mlm.__getitem__(padded_seq)
+            input_ids = torch.tensor([input_ids]).to(device)
+            logits = model(input_ids, idx_msa)
+            mask_logits = logits[0, mask_position, :]
+            predicted_token_probability, predicted_token_id = torch.topk((torch.softmax(mask_logits, dim=-1)), k=5)
+            topk.append(predicted_token_id)
+            predicted_token = vocab_mlm.idx_to_token[predicted_token_id[0].item()]
+            predicted_token_probability_all.append(predicted_token_probability[0].item())
+            padded_seq[mask_position] = predicted_token
+
+        cls_pos = vocab_mlm.to_tokens(list(topk[0]))
+        if X1 != "X":
+            Topk = X1
+            Subtype = X1
+            Potency = padded_seq[2],predicted_token_probability_all[0]
+        elif X2 != "X":
             Topk = cls_pos
-
+            Subtype = padded_seq[1],predicted_token_probability_all[0]
+            Potency = X2
         else:
-            probs_known = {}
-
-            for pos, token in enumerate(padded_seq):
-                if token in ["|", "[PAD]"]:
-                    continue
-
-                original_token = token
-                padded_seq[pos] = "[MASK]"
-
-                input_ids = torch.tensor(
-                    [vocab_mlm.__getitem__(padded_seq)]
-                ).to(device)
-
-                logits = model(input_ids, idx_msa)
-                mask_logits = logits[0, pos, :]
-                probs = torch.softmax(mask_logits, dim=-1)
-
-                token_id = vocab_mlm.token_to_idx[original_token]
-                probs_known[pos] = probs[token_id].item()
-
-                padded_seq[pos] = original_token
-
-            Subtype = (X1, probs_known.get(0, None))
-            Potency = (X2, probs_known.get(2, None))
-            Topk = "All known (no prediction)"
-
+            Topk = cls_pos
+            Subtype = padded_seq[1],predicted_token_probability_all[0]
+            Potency = padded_seq[2],predicted_token_probability_all[1]
     return Subtype, Potency, Topk
 
-
 iface = gr.Interface(
     fn=CTXGen,
     inputs=[