Update app.py

app.py

@@ -66,7 +66,7 @@ def CTXGen(τ, g_num, length_range, progress=gr.Progress()):
         start_time = time.time()
         while count < gen_num:
             if is_stopped:  # 检查是否停止
-                return 'output.csv', f"Generation stopped. Generated {count} conotoxins."
+                return pd.DataFrame(), f"Generation stopped. Generated {count} conotoxins."
 
             if time.time() - start_time > 1200:
                 break
@@ -84,7 +84,7 @@ def CTXGen(τ, g_num, length_range, progress=gr.Progress()):
 
             for i in range(length):
                 if is_stopped:  # 检查是否停止
-                    return 'output.csv', f"Generation stopped. Generated {count} conotoxins."
+                    return pd.DataFrame(), f"Generation stopped. Generated {count} conotoxins."
 
                 _, 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)
@@ -130,11 +130,19 @@ def CTXGen(τ, g_num, length_range, progress=gr.Progress()):
                     cls_probability_all.append(cls_probability)
                     act_pos_all.append(act_pos)
                     act_probability_all.append(act_probability)
-                    out = pd.DataFrame({'Generated_seq': generated_seqs_FINAL, 'Subtype': cls_pos_all, 'Subtype_probability': cls_probability_all, 'Potency': act_pos_all, 'Potency_probability': act_probability_all, 'random_seed': seed})
+                    out = pd.DataFrame({
+                        'Generated_seq': generated_seqs_FINAL,
+                        'Subtype': cls_pos_all,
+                        'Subtype_probability': cls_probability_all,
+                        'Potency': act_pos_all,
+                        'Potency_probability': act_probability_all,
+                        'random_seed': seed
+                    })
                     out.to_csv("output.csv", index=False, encoding='utf-8-sig')
                     count += 1
                     progress(count / gen_num, desc="Generating conotoxins...")
-    return 'output.csv', f"Generated {count} conotoxins."
+                    yield out, f"Generated {count} conotoxins."
+    return out, f"Generated {count} conotoxins."
 
 # 使用 gr.Blocks 构建界面
 with gr.Blocks() as demo:
@@ -147,11 +155,13 @@ with gr.Blocks() as demo:
         start_button = gr.Button("Start Generation")
         stop_button = gr.Button("Stop Generation")
     with gr.Row():
-        output_file = gr.File(label="Download generated conotoxins")
+        output_df = gr.DataFrame(label="Generated Conotoxins")
         status_text = gr.Textbox(label="Status")
+    with gr.Row():
+        output_file = gr.File(label="Download generated conotoxins")
     
     # 绑定事件
-    start_button.click(CTXGen, inputs=[τ, g_num, length_range], outputs=[output_file, status_text])
+    start_button.click(CTXGen, inputs=[τ, g_num, length_range], outputs=[output_df, status_text], live=True)
     stop_button.click(stop_generation, outputs=status_text)
 
 # 启动 Gradio 应用