Update app.py

app.py

@@ -4,16 +4,22 @@ import gradio as gr
 from gradio_rangeslider import RangeSlider
 import pandas as pd
 from utils import create_vocab, setup_seed
-from dataset_mlm import  get_paded_token_idx_gen, add_tokens_to_vocab
-import time 
-seed = random.randint(0,100000)
+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)
-save_path = 'mlm-model-27.pt'  
-train_seqs = pd.read_csv('C0_seq.csv')  
+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)
@@ -24,7 +30,14 @@ def temperature_sampling(logits, temperature):
     sampled_token = torch.multinomial(probabilities, 1)
     return sampled_token
 
+def stop_generation():
+    global is_stopped
+    is_stopped = True
+    return "Generation stopped."
+
 def CTXGen(τ, g_num, length_range, progress=gr.Progress()):
+    global is_stopped
+    is_stopped = False  # 重置停止标志
     start, end = length_range
     X1 = "X"
     X2 = "X"
@@ -43,29 +56,36 @@ def CTXGen(τ, g_num, length_range, progress=gr.Progress()):
 
         count = 0
         gen_num = int(g_num)
-        NON_AA = ["B","O","U","Z","X",'<K16>', '<α1β1γδ>', '<Ca22>', '<AChBP>', '<K13>', '<α1BAR>', '<α1β1ε>', '<α1AAR>', '<GluN3A>', '<α4β2>',
-                        '<GluN2B>', '<α75HT3>', '<Na14>', '<α7>', '<GluN2C>', '<NET>', '<NavBh>', '<α6β3β4>', '<Na11>', '<Ca13>', 
-                        '<Ca12>', '<Na16>', '<α6α3β2>', '<GluN2A>', '<GluN2D>', '<K17>', '<α1β1δε>', '<GABA>', '<α9>', '<K12>', 
-                        '<Kshaker>', '<α3β4>', '<Na18>', '<α3β2>', '<α6α3β2β3>', '<α1β1δ>', '<α6α3β4β3>', '<α2β2>','<α6β4>', '<α2β4>',
-                        '<Na13>', '<Na12>', '<Na15>', '<α4β4>', '<α7α6β2>', '<α1β1γ>', '<NaTTXR>', '<K11>', '<Ca23>', 
-                        '<α9α10>','<α6α3β4>', '<NaTTXS>', '<Na17>','<high>','<low>','[UNK]','[SEP]','[PAD]','[CLS]','[MASK]']
+        NON_AA = ["B", "O", "U", "Z", "X", '<K16>', '<α1β1γδ>', '<Ca22>', '<AChBP>', '<K13>', '<α1BAR>', '<α1β1ε>', '<α1AAR>', '<GluN3A>', '<α4β2>',
+                  '<GluN2B>', '<α75HT3>', '<Na14>', '<α7>', '<GluN2C>', '<NET>', '<NavBh>', '<α6β3β4>', '<Na11>', '<Ca13>',
+                  '<Ca12>', '<Na16>', '<α6α3β2>', '<GluN2A>', '<GluN2D>', '<K17>', '<α1β1δε>', '<GABA>', '<α9>', '<K12>',
+                  '<Kshaker>', '<α3β4>', '<Na18>', '<α3β2>', '<α6α3β2β3>', '<α1β1δ>', '<α6α3β4β3>', '<α2β2>', '<α6β4>', '<α2β4>',
+                  '<Na13>', '<Na12>', '<Na15>', '<α4β4>', '<α7α6β2>', '<α1β1γ>', '<NaTTXR>', '<K11>', '<Ca23>',
+                  '<α9α10>', '<α6α3β4>', '<NaTTXS>', '<Na17>', '<high>', '<low>', '[UNK]', '[SEP]', '[PAD]', '[CLS]', '[MASK]']
 
         start_time = time.time()
         while count < gen_num:
+            if is_stopped:  # 检查是否停止
+                return 'output.csv', f"Generation stopped. Generated {count} conotoxins."
+
             if time.time() - start_time > 1200:
                 break
+
             gen_len = random.randint(int(start), int(end))
             X3 = "X" * gen_len
             seq = [f"{X1}|{X2}|{X3}|{X4}|{X5}|{X6}"]
             vocab_mlm.token_to_idx["X"] = 4
 
             padded_seq, _, _, _ = get_paded_token_idx_gen(vocab_mlm, seq, new_seq)
-            input_text = ["[MASK]" if i=="X" else i for i in padded_seq]
+            input_text = ["[MASK]" if i == "X" else i for i in padded_seq]
 
             gen_length = len(input_text)
             length = gen_length - sum(1 for x in input_text if x != '[MASK]')
 
             for i in range(length):
+                if is_stopped:  # 检查是否停止
+                    return 'output.csv', 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)
                 idx_msa = torch.tensor(idx_msa).unsqueeze(0).to(device)
@@ -73,27 +93,26 @@ def CTXGen(τ, g_num, length_range, progress=gr.Progress()):
 
                 mask_positions = [j for j in range(gen_length) if input_text[j] == "[MASK]"]
                 mask_position = torch.tensor([mask_positions[torch.randint(len(mask_positions), (1,))]])
-                
-                logits = model(idx_seq,idx_msa, attn_idx) 
+
+                logits = model(idx_seq, idx_msa, attn_idx)
                 mask_logits = logits[0, mask_position.item(), :]
 
                 predicted_token_id = temperature_sampling(mask_logits, τ)
-
                 predicted_token = vocab_mlm.to_tokens(int(predicted_token_id))
                 input_text[mask_position.item()] = predicted_token
                 padded_seq[mask_position.item()] = predicted_token.strip()
                 new_seq = padded_seq
 
             generated_seq = input_text
-        
+
             generated_seq[1] = "[MASK]"
             generated_seq[2] = "[MASK]"
             input_ids = vocab_mlm.__getitem__(generated_seq)
             logits = model(torch.tensor([input_ids]).to(device), idx_msa)
-            
+
             cls_mask_logits = logits[0, 1, :]
             act_mask_logits = logits[0, 2, :]
-            
+
             cls_probability, cls_mask_probs = torch.topk((torch.softmax(cls_mask_logits, dim=-1)), k=1)
             act_probability, act_mask_probs = torch.topk((torch.softmax(act_mask_logits, dim=-1)), k=1)
 
@@ -117,16 +136,24 @@ def CTXGen(τ, g_num, length_range, progress=gr.Progress()):
                     progress(count / gen_num, desc="Generating conotoxins...")
     return 'output.csv', f"Generated {count} conotoxins."
 
+# 定义 Gradio 界面
 iface = gr.Interface(
     fn=CTXGen,
     inputs=[
         gr.Slider(minimum=1, maximum=2, step=0.1, label="τ"),
-        gr.Dropdown(choices=[1,10,100], label="Number of generations"),
+        gr.Dropdown(choices=[1, 10, 100], label="Number of generations"),
         RangeSlider(minimum=8, maximum=50, step=1, value=(12, 17), label="Length range")
     ],
     outputs=[
         gr.File(label="Download generated conotoxins"),
         gr.Textbox(label="Progress")
-    ]
+    ],
+    live=True
 )
+
+# 添加停止按钮
+stop_button = gr.Button("Stop Generation")
+stop_button.click(stop_generation, outputs=gr.Textbox(label="Status"))
+
+# 启动 Gradio 应用
 iface.launch()