Update app.py

app.py

@@ -85,7 +85,14 @@ def CTXGen(X0, X1, X2, τ, g_num):
         cls_proba_parent = cls_probability_parent[cls_pos_parent.index(X1)].item()
         act_proba_parent = act_probability_parent[act_pos_parent.index(X2)].item()
 
+        start_time = time.time()
         while count < gen_num:
+            if is_stopped:
+                return pd.DataFrame(), "output.csv"
+
+            if time.time() - start_time > 1200:
+                break
+
             gen_len = len(X0)
             seq = [f"{X1}|{X2}|{X3}|{X4}|{X5}|{X6}"]
             vocab_mlm.token_to_idx["X"] = 4
@@ -97,6 +104,9 @@ def CTXGen(X0, X1, X2, τ, g_num):
             length = gen_length - sum(1 for x in input_text if x != '[MASK]')
 
             for i in range(length):
+                if is_stopped:
+                    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)
                 idx_msa = torch.tensor(idx_msa).unsqueeze(0).to(device)