gen

app.py

@@ -3,6 +3,7 @@ from transformers import AutoTokenizer, AutoModelForMaskedLM
 import torch
 from torch.distributions.categorical import Categorical
 import numpy as np
+import pandas as pd
 
 # Load the model and tokenizer
 tokenizer = AutoTokenizer.from_pretrained("TianlaiChen/PepMLM-650M")
@@ -57,8 +58,17 @@ def generate_peptide(protein_seq, peptide_length, top_k, num_binders):
 
         # Add the generated binder and its PPL to the results list
         binders_with_ppl.append([generated_binder, ppl_value])
+        
+        # Convert the list of lists to a pandas dataframe
+        df = pd.DataFrame(binders_with_ppl, columns=["Binder", "Perplexity"])
+    
+        # Save the dataframe to a CSV file
+        output_filename = "output.csv"
+        df.to_csv(output_filename, index=False)
+
+
+    return binders_with_ppl, output_filename
 
-    return binders_with_ppl
 
 # Define the Gradio interface
 interface = gr.Interface(
@@ -67,13 +77,16 @@ interface = gr.Interface(
         gr.Textbox(label="Protein Sequence", info="Enter protein sequence here", type="text"),
         gr.Slider(3, 50, value=15, label="Peptide Length", step=1, info='Default value is 15'),
         gr.Slider(1, 10, value=3, label="Top K Value", step=1, info='Default value is 3'),
-        gr.Dropdown(choices=[1, 2, 4, 8, 16, 32], label="Number of Binders", default=1)
+        gr.Dropdown(choices=[1, 2, 4, 8, 16, 32], label="Number of Binders", value=1)
     ],
-    outputs=gr.Dataframe(
+    outputs=[
+    gr.Dataframe(
         headers=["Binder", "Perplexity"],
         datatype=["str", "number"],
         col_count=(2, "fixed")
     ),
+    gr.outputs.File(label="Download CSV")
+    ],
     title="PepMLM: Target Sequence-Conditioned Generation of Peptide Binders via Masked Language Modeling"
 )
 

requirements.txt

@@ -1,4 +1,5 @@
 transformers
 gradio
 torch
-numpy
+numpy
+pandas