Update app.py

app.py

@@ -7,14 +7,11 @@ from utils import create_vocab, setup_seed
 from dataset_mlm import get_paded_token_idx_gen, add_tokens_to_vocab
 import time
 
-# 全局标志，用于控制停止
 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)
@@ -37,7 +34,7 @@ def stop_generation():
 
 def CTXGen(τ, g_num, length_range):
     global is_stopped
-    is_stopped = False  # 重置停止标志
+    is_stopped = False
     start, end = length_range
     X1 = "X"
     X2 = "X"
@@ -65,7 +62,7 @@ def CTXGen(τ, g_num, length_range):
 
         start_time = time.time()
         while count < gen_num:
-            if is_stopped:  # 检查是否停止
+            if is_stopped:
                 return pd.DataFrame(), "output.csv"
 
             if time.time() - start_time > 1200:
@@ -83,7 +80,7 @@ def CTXGen(τ, g_num, length_range):
             length = gen_length - sum(1 for x in input_text if x != '[MASK]')
 
             for i in range(length):
-                if is_stopped:  # 检查是否停止
+                if is_stopped:
                     return pd.DataFrame(), "output.csv"
 
                 _, idx_seq, idx_msa, attn_idx = get_paded_token_idx_gen(vocab_mlm, seq, new_seq)
@@ -143,13 +140,12 @@ def CTXGen(τ, g_num, length_range):
                     yield out, "output.csv"
     return out, "output.csv"
 
-# 使用 gr.Blocks 构建界面
 with gr.Blocks() as demo:
     gr.Markdown("# Conotoxin Generation")
     with gr.Row():
         τ = gr.Slider(minimum=1, maximum=2, step=0.1, label="τ")
-        g_num = gr.Dropdown(choices=[1, 10, 100], label="Number of generations")
-        length_range = RangeSlider(minimum=8, maximum=50, step=1, value=(12, 17), label="Length range")
+        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")
     with gr.Row():
         start_button = gr.Button("Start Generation")
         stop_button = gr.Button("Stop Generation")
@@ -158,9 +154,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])
     stop_button.click(stop_generation, outputs=None)
 
-# 启动 Gradio 应用
 demo.launch()