gen

app.py

@@ -3,14 +3,14 @@ from transformers import AutoTokenizer, AutoModelForMaskedLM
 import torch
 from torch.distributions.categorical import Categorical
 
-'''
+
 # Load the model and tokenizer
 tokenizer = AutoTokenizer.from_pretrained("TianlaiChen/PepMLM-650M")
 model = AutoModelForMaskedLM.from_pretrained("TianlaiChen/PepMLM-650M")
 
 
 def generate_peptide(protein_seq, peptide_length, top_k):
-    print('Within')
+
     peptide_length = int(peptide_length)
     top_k = int(top_k)
     
@@ -31,19 +31,18 @@ def generate_peptide(protein_seq, peptide_length, top_k):
 
     generated_peptide = tokenizer.decode(predicted_token_ids, skip_special_tokens=True)
     return f"Generated Sequence: {generated_peptide.replace(' ', '')}"
-'''
 
-def hello():
-    print('hello')
-    return "hello"
+
 
 # Define the Gradio interface
 interface = gr.Interface(
-    fn=hello,
+    fn=generate_peptide,
     inputs=[
-        gr.Textbox(label="Protein Sequence", default="Enter protein sequence here", type="text"),
-        gr.Dropdown(choices=[str(i) for i in range(2, 51)], label="Peptide Length", default="15"),
-        gr.Dropdown(choices=[str(i) for i in range(1, 11)], label="Top K Value", default="3")
+        gr.Textbox(label="Protein Sequence", info = "Enter protein sequence here", type="text"),
+        gr.Dropdown(choices=[str(i) for i in range(2, 51)], label="Peptide Length", default="15",
+        info='Default value is 15'),
+        gr.Dropdown(choices=[str(i) for i in range(1, 11)], label="Top K Value", default="3",
+        info='Default value is 3')
     ],
     outputs="textbox",
 )