Update app.py

app.py

@@ -66,7 +66,7 @@ def CTXGen(τ, g_num, length_range):
         start_time = time.time()
         while count < gen_num:
             if is_stopped:  # 检查是否停止
-                yield pd.DataFrame(), f"Generation stopped. Generated {count} conotoxins."
+                return pd.DataFrame(), "output.csv"
 
             if time.time() - start_time > 1200:
                 break
@@ -84,7 +84,7 @@ def CTXGen(τ, g_num, length_range):
 
             for i in range(length):
                 if is_stopped:  # 检查是否停止
-                    yield pd.DataFrame(), f"Generation stopped. Generated {count} conotoxins."
+                    return pd.DataFrame(), "output.csv"
 
                 _, idx_seq, idx_msa, attn_idx = get_paded_token_idx_gen(vocab_mlm, seq, new_seq)
                 idx_seq = torch.tensor(idx_seq).unsqueeze(0).to(device)
@@ -140,11 +140,12 @@ def CTXGen(τ, g_num, length_range):
                     })
                     out.to_csv("output.csv", index=False, encoding='utf-8-sig')
                     count += 1
-                    yield out
-    yield out
+                    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")
@@ -158,8 +159,8 @@ with gr.Blocks() as demo:
         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=output_file)
+    start_button.click(CTXGen, inputs=[τ, g_num, length_range], outputs=[output_df, output_file])
+    stop_button.click(stop_generation, outputs=None)
 
 # 启动 Gradio 应用
 demo.launch()