Update app.py

app.py

@@ -6,12 +6,9 @@ from dataset_mlm import  get_paded_token_idx_gen, add_tokens_to_vocab
 import gradio as gr
 from gradio_rangeslider import RangeSlider
 import time
-import numba
-from numba import objmode
 
 is_stopped = False
 
-@numba.jit(nopython=True)
 def temperature_sampling(logits, temperature):
     logits = logits / temperature
     probabilities = torch.softmax(logits, dim=-1)
@@ -23,29 +20,6 @@ def stop_generation():
     is_stopped = True
     return "Generation stopped."
 
-@numba.jit(nopython=False)
-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,))]])
-        
-        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
-    return input_text
-
 def CTXGen(X0, X3, X1, X2, τ, g_num, model_name, seed):
     if seed =='random':
         seed = random.randint(0,100000)
@@ -131,11 +105,31 @@ 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(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,))]])
+                
+                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]"
             input_ids = vocab_mlm.__getitem__(generated_seq)