Update app.py

app.py

@@ -8,7 +8,8 @@ import matplotlib
 
 matplotlib.use("Agg")
 import matplotlib.pyplot as plt
-from transformers import pipeline as pl
+#from transformers import pipeline as pl
+from transformers import GPT2LMHeadModel , GPT2Tokenizer
 from GPUtil import showUtilization as gpu_usage
 
 import pandas as pd
@@ -20,6 +21,7 @@ import torch
 import gc
 import jax
 from numba import cuda
+import math
 print('GPU available',torch.cuda.is_available())
 #print('__CUDA Device Name:',torch.cuda.get_device_name(0))
 print(os.getcwd())
@@ -81,27 +83,58 @@ def predict_structure(prefix, feature_dict, model_runners, random_seed=0):
             f.write(protein.to_pdb(unrelaxed_protein))
     return plddts
 
+def compute_perplexity(model, tokenizer, sequence):
+    input_ids = torch.tensor(tokenizer.encode(sentence)).unsqueeze(0)
+    with torch.no_grad():
+        outputs = model(input_ids, labels=input_ids)
+    loss, logits = outputs[:2]
+    return math.exp(loss)
+
 @ray.remote(num_gpus=1, max_calls=1)
 def run_protgpt2(startsequence, length, repetitionPenalty, top_k_poolsize, max_seqs):
     print("running protgpt2")
     print(gpu_usage())
-    protgpt2 = pl("text-generation", model="nferruz/ProtGPT2")
-    sequences = protgpt2(
-        startsequence,
-        max_length=length,
-        do_sample=True,
-        top_k=top_k_poolsize,
-        repetition_penalty=repetitionPenalty,
-        num_return_sequences=max_seqs,
-        eos_token_id=0,
-    )
+    seqs_to_sample = max_seqs*10 # get the top 10
+    #protgpt2 = pl("text-generation", model="nferruz/ProtGPT2")
+    model = GPT2LMHeadModel.from_pretrained("nferruz/ProtGPT2")
+    tokenizer = GPT2Tokenizer.from_pretrained("nferruz/ProtGPT2")
+    input_ids = tokenizer.encode(startsequence, return_tensors='pt')
+    
+    sequences = model.generate(input_ids,
+                    max_length=length, 
+                    do_sample=True, 
+                    top_k=top_k_poolsize,  
+                    repetition_penalty=repetitionPenalty, 
+                    num_return_sequences=seqs_to_sample,
+                    eos_token_id=0)
+    filtered_sequences = []
+    for sequence in sequences:
+        decoded_seq = tokenizer.decode(seq)
+        # No newlines in first line and avoid truncation
+        if '\n' not in decoded_seq[0:60] and decoded_seq.count('<|endoftext|>')>=2:
+            clean_seq = decoded_seq.split('<|endoftext|>')[0]
+            ppl = compute_perplexity(model, tokenizer, clean_seq)
+            filtered_sequences.append((clean_seq, ppl/len(clean_seq)))
+            
+    ## THis needs to be fixed to show warning if not enough sequences fulfill the criteria!         
+    selected_sequences = filtered_sequences.sort(key = lambda x: x[2])[:max_seqs]     
+#    sequences = protgpt2(
+#        startsequence,
+#        max_length=length,
+#        do_sample=True,
+#        top_k=top_k_poolsize,
+#        repetition_penalty=repetitionPenalty,
+#        num_return_sequences=seqs_to_sample,
+#        eos_token_id=0,
+#    )
+
     print("Cleaning up after protGPT2")
     #print(gpu_usage())
     #torch.cuda.empty_cache()
     #device = cuda.get_current_device()
     #device.reset()
     #print(gpu_usage())
-    return sequences
+    return selected_sequences
     
 @ray.remote(num_gpus=1, max_calls=1)
 def run_alphafold(startsequence):
@@ -140,9 +173,13 @@ def update_protGPT2(inp, length,repetitionPenalty, top_k_poolsize, max_seqs):
     seqlen = length
     generated_seqs = ray.get(run_protgpt2.remote(startsequence, seqlen, repetitionPenalty, top_k_poolsize, max_seqs))
     gen_seqs = [x["generated_text"] for x in generated_seqs]
-    print(gen_seqs)
+    # Make sure sequences weren't truncated due to the length cutoff
+    
+    # Select the best scoring top 10th:
+    
+    print(sel_seqs)
     sequencestxt = ""
-    for i, seq in enumerate(gen_seqs):
+    for i, seq in enumerate(sel_seqs):
         s = seq.replace("\n","")
         seqlen = len(s)
         s = "\n".join([s[i:i+70] for i in range(0, len(s), 70)])
@@ -360,8 +397,8 @@ with proteindream:
     )
     with gr.Box():
         with gr.Row():
-            inp = gr.Textbox(placeholder="M", label="Start sequence")
-            length = gr.Number(value=50, label="Max sequence length")
+            inp = gr.Textbox(placeholder="MTYKLILNGKTLKGETTT", label="Start sequence")
+            length = gr.Number(value=100, label="Max sequence length")
         with gr.Row():
             repetitionPenalty = gr.Slider(minimum=1, maximum=5,value=1.2, label="Repetition penalty")
             top_k_poolsize = gr.Slider(minimum=700, maximum=52056,value=950, label="Top-K sampling pool size")