Update app.py

app.py

@@ -6,9 +6,11 @@ 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
 
 is_stopped = False
 
+@numba.jit(nopython=True)
 def temperature_sampling(logits, temperature):
     logits = logits / temperature
     probabilities = torch.softmax(logits, dim=-1)
@@ -20,6 +22,24 @@ def stop_generation():
     is_stopped = True
     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]')
+    for i in range(length):
+        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)
@@ -106,30 +126,7 @@ 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]')
-            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 = generate_sequence(input_text, model, vocab_mlm, idx_msa, τ)
         
             generated_seq[1] = "[MASK]"
             input_ids = vocab_mlm.__getitem__(generated_seq)