Update app.py

app.py

@@ -79,17 +79,18 @@ def predict_structure(prefix, feature_dict, model_runners, random_seed=0):
     return plddts
 
 
-def run_protgpt2(startsequence, length):
+def run_protgpt2(startsequence, length, repetitionPenalty, top_k_poolsize, max_seqs):
     protgpt2 = pl("text-generation", model="nferruz/ProtGPT2")
     sequences = protgpt2(
         startsequence,
         max_length=length,
         do_sample=True,
-        top_k=950,
-        repetition_penalty=1.2,
-        num_return_sequences=5,
+        top_k=top_k_poolsize,
+        repetition_penalty=repetitionPenalty,
+        num_return_sequences=max_seqs,
         eos_token_id=0,
     )
+    del protgpt2
     torch.cuda.empty_cache()
     return sequences
 
@@ -342,11 +343,15 @@ with proteindream:
     with gr.Group():
         with gr.Row():
             inp = gr.Textbox(placeholder="M", label="Start sequence")
-            length = gr.Number(value=50, label="Target sequence length")
+            length = gr.Number(value=50, 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")
+            max_seqs = gr.Slider(minimum=2, maximum=20,value=5, label="Number of sequences to generate")
         btn = gr.Button("Predict sequences using protGPT2")
 
     results = gr.Textbox(label="Results", lines=15)
-    btn.click(fn=update_protGPT2, inputs=[inp, length], outputs=results)
+    btn.click(fn=update_protGPT2, inputs=[inp, length, repetitionPenalty, top_k_poolsize, max_seqs], outputs=results)
 
     gr.Markdown("## AlphaFold")
     gr.Markdown(