embedings generate

.gitignore

@@ -1,3 +1,4 @@
 .DS_Store
 .idea/
 ps4_data/__pycache__/
+ps4_eval/__pycache__/

app.py

@@ -1,12 +1,12 @@
 import gradio as gr
-from ps4_models.classifiers import *
+from ps4_eval.eval import sample_new_sequence
 from ps4_data.get_embeddings import generate_embedings
 
 
 def pred(residue_seq):
-    generate_embedings(residue_seq)
-    model = PS4_Mega()
-    return "Hello " + residue_seq + "!!"
+    embs = generate_embedings(residue_seq)["residue_embs"]["0"]
+    preds = sample_new_sequence(embs, "ps4_models/Mega/PS4-Mega_loss-0.633_acc-78.176.pt")
+    return preds
 
 
 iface = gr.Interface(fn=pred, title="Protein Secondary Structure Prediction with PS4-Mega",

ps4_data/data/protT5/output/per_residue_embeddings0.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ce32f0bb1cced643cc12d813b19bfa11b074b0d789329d7da1a1c066d63ccf75
+size 197778

ps4_data/get_embeddings.py

@@ -22,17 +22,10 @@ def generate_embedings(input_seq, output_path=None):
 
     # Load fasta.
     all_seqs = {"0": input_seq}
-
-    chunk_size = 1000
-
     # Compute embeddings and/or secondary structure predictions
-    for i in range(0, len(all_seqs), chunk_size):
-        keys = list(all_seqs.keys())[i: chunk_size + i]
-        seqs = {k: all_seqs[k] for k in keys}
-        results = __get_embeddings(model, tokenizer, seqs, device)
+    results = __get_embeddings(model, tokenizer, all_seqs, device)
 
-        # Store per-residue embeddings
-        __save_embeddings(results["residue_embs"], per_residue_path + f"{i}.npz")
+    return results
 
 
 def __get_T5_model(device):
@@ -92,7 +85,7 @@ def __get_embeddings(model, tokenizer, seqs, device, per_residue=True,
                 # slice off padding --> batch-size x seq_len x embedding_dim
                 emb = embedding_repr.last_hidden_state[batch_idx, :s_len]
                 if per_residue:  # store per-residue embeddings (Lx1024)
-                    results["residue_embs"][identifier] = emb.detach().cpu().numpy().squeeze()
+                    results["residue_embs"][identifier] = emb.detach().cpu().squeeze()
                     print("emb_count:", len(results["residue_embs"]))
 
     passed_time = time.time() - start

ps4_eval/eval.py

@@ -0,0 +1,67 @@
+import torch
+from torch import nn, cuda, load, device
+from ps4_models.classifiers import PS4_Mega, PS4_Conv
+
+
+def load_trained_model(load_path, model_name='PS4_Mega'):
+
+    if model_name.lower() not in ['ps4_conv', 'ps4_mega']:
+        raise ValueError(f'Model name {model_name} not recognised, please choose from PS4_Conv, PS4_Mega')
+
+    model: nn.Module = PS4_Mega() if model_name.lower() == 'ps4_mega' else PS4_Conv()
+
+    if load_path != '':
+        try:
+            if cuda.is_available():
+                model.load_state_dict(load(load_path)['model_state_dict'])
+            else:
+                model.load_state_dict(load(load_path, map_location=device('cpu'))['model_state_dict'])
+            print("loded params from", load_path)
+        except:
+            raise ImportError(f'No file located at {load_path}, could not load parameters')
+    print(model)
+
+    pytorch_total_params = sum(par.numel() for par in model.parameters() if par.requires_grad)
+    print(pytorch_total_params)
+
+    return model
+
+
+# MARK: sampling from new sequence
+def sample_new_sequence(embs, weights_load_path, model_name='PS4_Mega'):
+
+    model = load_trained_model(weights_load_path, model_name)
+
+    seq_size = len(embs)
+    R = embs.view(1, seq_size)
+
+    pred_ss = ''
+
+    with torch.no_grad():
+        y_hat = model(R)
+        probs = torch.softmax(y_hat, 2)
+        _, ss_preds = torch.max(probs, 2)
+
+        for i in range(seq_size):
+            ss = ss_preds[0][i].item()
+            ss = ss_tokeniser(ss, reverse=True)
+            pred_ss += ss
+
+    return pred_ss
+
+
+def ss_tokeniser(ss, reverse=False):
+
+    ss_set = ['C', 'T', 'G', 'H', 'S', 'B', 'I', 'E', 'C']
+
+    if reverse:
+        return inverse_ss_tokeniser(ss)
+    else:
+        return 0 if (ss == 'P' or ss == ' ') else ss_set.index(ss)
+
+
+def inverse_ss_tokeniser(ss):
+
+    ss_set = ['C', 'T', 'G', 'H', 'S', 'B', 'I', 'E', 'C', 'C']
+
+    return ss_set[ss]

requirements.txt

@@ -5,5 +5,5 @@ scikit-learn~=0.24.2
 transformers~=4.26.1
 setuptools~=57.4.0
 pandas~=1.3.2
-wget~=3.2
+sentencepiece~=0.1.97
 -e git+https://github.com/facebookresearch/mega.git@main#egg=fairseq