Update app.py

app.py

@@ -23,10 +23,15 @@ def stop_generation():
     return "Generation stopped."
 
 @numba.jit(nopython=True)
-def generate_sequence(input_text, model, vocab_mlm, idx_msa, τ):
-    gen_length = len(input_text)
-    length = gen_length - sum(1 for x in input_text if x != '[MASK]')
+def generate_sequence(length, vocab_mlm, seq, new_seq, τ, input_text):
     for i in range(length):
+        if is_stopped:
+            return "output.csv", pd.DataFrame()
+        
+        _, 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)
+        attn_idx = torch.tensor(attn_idx).to(device)
         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,))]])
         
@@ -125,8 +130,11 @@ def CTXGen(X0, X3, X1, X2, τ, g_num, model_name, seed):
 
             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]
+            gen_length = len(input_text)
+            length = gen_length - sum(1 for x in input_text if x != '[MASK]')
 
-            generated_seq = generate_sequence(input_text, model, vocab_mlm, idx_msa, τ)
+            #函数
+            generated_seq = generate_sequence(length, vocab_mlm, seq, new_seq, τ, input_text)
         
             generated_seq[1] = "[MASK]"
             input_ids = vocab_mlm.__getitem__(generated_seq)