Update app.py

app.py

@@ -12,15 +12,6 @@ is_stopped = False
 seed = random.randint(0, 100000)
 setup_seed(seed)
 
-device = torch.device("cpu")
-vocab_mlm = create_vocab()
-vocab_mlm = add_tokens_to_vocab(vocab_mlm)
-save_path = 'mlm-model-27.pt'
-train_seqs = pd.read_csv('C0_seq.csv')
-train_seq = train_seqs['Seq'].tolist()
-model = torch.load(save_path, map_location=torch.device('cpu'))
-model = model.to(device)
-
 def temperature_sampling(logits, temperature):
     logits = logits / temperature
     probabilities = torch.softmax(logits, dim=-1)
@@ -32,9 +23,19 @@ def stop_generation():
     is_stopped = True
     return "Generation stopped."
 
-def CTXGen(τ, g_num, length_range):
+def CTXGen(τ, g_num, length_range, model_name):
     global is_stopped
     is_stopped = False
+
+    device = torch.device("cpu")
+    vocab_mlm = create_vocab()
+    vocab_mlm = add_tokens_to_vocab(vocab_mlm)
+    save_path = model_name
+    train_seqs = pd.read_csv('C0_seq.csv')
+    train_seq = train_seqs['Seq'].tolist()
+    model = torch.load(save_path, map_location=torch.device('cpu'))
+    model = model.to(device)
+
     start, end = length_range
     X1 = "X"
     X2 = "X"
@@ -43,7 +44,6 @@ def CTXGen(τ, g_num, length_range):
     X6 = ""
     model.eval()
     with torch.no_grad():
-        new_seq = None
         IDs = []
         generated_seqs = []
         generated_seqs_FINAL = []
@@ -63,6 +63,7 @@ def CTXGen(τ, g_num, length_range):
 
         start_time = time.time()
         while count < gen_num:
+            new_seq = None
             if is_stopped:
                 return pd.DataFrame(), "output.csv"
 
@@ -149,6 +150,7 @@ with gr.Blocks() as demo:
         τ = gr.Slider(minimum=1, maximum=2, step=0.1, label="τ")
         g_num = gr.Dropdown(choices=[1, 10, 20, 30, 40, 50], label="Number of generations")
         length_range = RangeSlider(minimum=8, maximum=50, step=1, value=(12, 16), label="Length range")
+        model_name = gr.Dropdown(choices=['model_final.pt','model_C1.pt','model_C2.pt','model_C3.pt','model_C4.pt','model_C5.pt','model_mlm.pt'], label="Model")
     with gr.Row():
         start_button = gr.Button("Start Generation")
         stop_button = gr.Button("Stop Generation")
@@ -157,7 +159,7 @@ with gr.Blocks() as demo:
     with gr.Row():
         output_file = gr.File(label="Download generated conotoxins")
     
-    start_button.click(CTXGen, inputs=[τ, g_num, length_range], outputs=[output_df, output_file])
+    start_button.click(CTXGen, inputs=[τ, g_num, length_range, model_name], outputs=[output_df, output_file])
     stop_button.click(stop_generation, outputs=None)
 
 demo.launch()