File size: 24,153 Bytes
6bd3a1e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
import streamlit as st
from io import StringIO
from Bio import SeqIO

import numpy as np
import os
import pandas as pd
import random
import torch
import torch.nn as nn
import torch.nn.functional as F
from collections import Counter, OrderedDict
from copy import deepcopy
from esm import Alphabet, FastaBatchedDataset, ProteinBertModel, pretrained, MSATransformer
from esm.data import *
from esm.model.esm2 import ESM2
from torch import nn
from torch.nn import Linear
from torch.nn.utils.rnn import pad_sequence
from torch.utils.data import Dataset, DataLoader
seed = 19961231
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)


st.title("IRES-LM prediction and mutation")

# Input sequence
st.subheader("Input sequence")

seq = st.text_area("FASTA format only", value=">vir_CVB3_ires_00505.1\nTTAAAACAGCCTGTGGGTTGATCCCACCCACAGGCCCATTGGGCGCTAGCACTCTGGTATCACGGTACCTTTGTGCGCCTGTTTTATACCCCCTCCCCCAACTGTAACTTAGAAGTAACACACACCGATCAACAGTCAGCGTGGCACACCAGCCACGTTTTGATCAAGCACTTCTGTTACCCCGGACTGAGTATCAATAGACTGCTCACGCGGTTGAAGGAGAAAGCGTTCGTTATCCGGCCAACTACTTCGAAAAACCTAGTAACACCGTGGAAGTTGCAGAGTGTTTCGCTCAGCACTACCCCAGTGTAGATCAGGTCGATGAGTCACCGCATTCCCCACGGGCGACCGTGGCGGTGGCTGCGTTGGCGGCCTGCCCATGGGGAAACCCATGGGACGCTCTAATACAGACATGGTGCGAAGAGTCTATTGAGCTAGTTGGTAGTCCTCCGGCCCCTGAATGCGGCTAATCCTAACTGCGGAGCACACACCCTCAAGCCAGAGGGCAGTGTGTCGTAACGGGCAACTCTGCAGCGGAACCGACTACTTTGGGTGTCCGTGTTTCATTTTATTCCTATACTGGCTGCTTATGGTGACAATTGAGAGATCGTTACCATATAGCTATTGGATTGGCCATCCGGTGACTAATAGAGCTATTATATATCCCTTTGTTGGGTTTATACCACTTAGCTTGAAAGAGGTTAAAACATTACAATTCATTGTTAAGTTGAATACAGCAAA")
st.subheader("Upload sequence file")
uploaded = st.file_uploader("Sequence file in FASTA format")

# augments
global output_filename, start_nt_position, end_nt_position, mut_by_prob, transform_type, mlm_tok_num, n_mut, n_designs_ep, n_sampling_designs_ep, n_mlm_recovery_sampling, mutate2stronger
output_filename = st.text_input("output a .csv file", value='IRES_LM_prediction_mutation')
start_nt_position = st.number_input("The start position of the mutation of this sequence, the first position is defined as 0", value=0)
end_nt_position = st.number_input("The last position of the mutation of this sequence, the last position is defined as length(sequence)-1 or -1", value=-1)
mut_by_prob = st.checkbox("Mutated by predicted Probability or Transformed Probability of the sequence", value=True)
transform_type = st.selectbox("Type of probability transformation", 
                              ['', 'sigmoid', 'logit', 'power_law', 'tanh'],
                              index=2)
mlm_tok_num = st.number_input("Number of masked tokens for each sequence per epoch", value=1)
n_mut = st.number_input("Maximum number of mutations for each sequence", value=3)
n_designs_ep = st.number_input("Number of mutations per epoch", value=10)
n_sampling_designs_ep = st.number_input("Number of sampling mutations from n_designs_ep per epoch", value=5)
n_mlm_recovery_sampling = st.number_input("Number of MLM recovery samplings (with AGCT recovery)", value=1)
mutate2stronger = st.checkbox("Mutate to stronger IRES variant, otherwise mutate to weaker IRES", value=True)
if not mut_by_prob and transform_type != '':
    st.write("--transform_type must be '' when --mut_by_prob is False")
    transform_type = ''

global idx_to_tok, prefix, epochs, layers, heads, fc_node, dropout_prob, embed_dim, batch_toks, repr_layers, evaluation, include, truncate, return_contacts, return_representation, mask_toks_id, finetune

epochs = 5
layers = 6
heads = 16
embed_dim = 128
batch_toks = 4096
fc_node = 64
dropout_prob = 0.5
folds = 10
repr_layers = [-1]
include = ["mean"]
truncate = True
finetune = False
return_contacts = False
return_representation = False

global tok_to_idx, idx_to_tok, mask_toks_id
alphabet = Alphabet(mask_prob = 0.15, standard_toks = 'AGCT')
assert alphabet.tok_to_idx == {'<pad>': 0, '<eos>': 1, '<unk>': 2, 'A': 3, 'G': 4, 'C': 5, 'T': 6, '<cls>': 7, '<mask>': 8, '<sep>': 9}

# tok_to_idx = {'<pad>': 0, '<eos>': 1, '<unk>': 2, 'A': 3, 'G': 4, 'C': 5, 'T': 6, '<cls>': 7, '<mask>': 8, '<sep>': 9}
tok_to_idx = {'-': 0, '&': 1, '?': 2, 'A': 3, 'G': 4, 'C': 5, 'T': 6, '!': 7, '*': 8, '|': 9}
idx_to_tok = {idx: tok for tok, idx in tok_to_idx.items()}
# st.write(tok_to_idx)
mask_toks_id = 8

global w1, w2, w3
w1, w2, w3 = 1, 1, 100

class CNN_linear(nn.Module):
    def __init__(self):
        super(CNN_linear, self).__init__()
        
        self.esm2 = ESM2(num_layers = layers,
                         embed_dim = embed_dim,
                         attention_heads = heads,
                         alphabet = alphabet)
        
        self.dropout = nn.Dropout(dropout_prob)
        self.relu = nn.ReLU()
        self.flatten = nn.Flatten()
        self.fc = nn.Linear(in_features = embed_dim, out_features = fc_node)
        self.output = nn.Linear(in_features = fc_node, out_features = 2)        
            
    def predict(self, tokens):
        
        x = self.esm2(tokens, [layers], need_head_weights=False, return_contacts=False, return_representation = True)
        x_cls = x["representations"][layers][:, 0]
        
        o = self.fc(x_cls)
        o = self.relu(o)
        o = self.dropout(o)
        o = self.output(o)
        
        y_prob = torch.softmax(o, dim = 1)
        y_pred = torch.argmax(y_prob, dim = 1)
        
        if transform_type:
            y_prob_transformed = prob_transform(y_prob[:,1])
            return y_prob[:,1], y_pred, x['logits'], y_prob_transformed
        else:
            return y_prob[:,1], y_pred, x['logits'], o[:,1]
    
    def forward(self, x1, x2):
        logit_1, repr_1 = self.predict(x1)
        logit_2, repr_2 = self.predict(x2)
        return (logit_1, logit_2), (repr_1, repr_2)


def prob_transform(prob, **kwargs): # Logits
    """
    Transforms probability values based on the specified method.

    :param prob: torch.Tensor, the input probabilities to be transformed
    :param transform_type: str, the type of transformation to be applied
    :param kwargs: additional parameters for transformations
    :return: torch.Tensor, transformed probabilities
    """

    if transform_type == 'sigmoid':
        x0 = kwget('x0', 0.5)
        k = kwget('k', 10.0)
        prob_transformed = 1 / (1 + torch.exp(-k * (prob - x0)))
    
    elif transform_type == 'logit':
        # Adding a small value to avoid log(0) and log(1)
        prob_transformed = torch.log(prob + 1e-6) - torch.log(1 - prob + 1e-6)
    
    elif transform_type == 'power_law':
        gamma = kwget('gamma', 2.0)
        prob_transformed = torch.pow(prob, gamma)
    
    elif transform_type == 'tanh':
        k = kwget('k', 2.0)
        prob_transformed = torch.tanh(k * prob)
    
    return prob_transformed

def random_replace(sequence, continuous_replace=False):
    global start_nt_position, end_nt_position
    if end_nt_position == -1: end_nt_position = len(sequence)-1
    if start_nt_position < 0 or end_nt_position > len(sequence)-1 or start_nt_position > end_nt_position:
#         raise ValueError("Invalid start/end positions")
        st.write("Invalid start/end positions")
        start_nt_position, end_nt_position = 0, len(sequence)-1
    
    # 将序列切片成三部分:替换区域前、替换区域、替换区域后
    pre_segment = sequence[:start_nt_position]
    target_segment = list(sequence[start_nt_position:end_nt_position + 1])  # +1因为Python的切片是右开区间
    post_segment = sequence[end_nt_position + 1:]
    
    if not continuous_replace:
        # 随机替换目标片段的mlm_tok_num个位置
        indices = random.sample(range(len(target_segment)), mlm_tok_num)
        for idx in indices:
            target_segment[idx] = '*'
    else:
        # 在目标片段连续替换mlm_tok_num个位置
        max_start_idx = len(target_segment) - mlm_tok_num  # 确保从i开始的n_mut个元素不会超出目标片段的长度
        if max_start_idx < 1:  # 如果目标片段长度小于mlm_tok_num,返回原始序列
            return target_segment
        start_idx = random.randint(0, max_start_idx)
        for idx in range(start_idx, start_idx + mlm_tok_num):
            target_segment[idx] = '*'
    
    # 合并并返回最终的序列
    return ''.join([pre_segment] + target_segment + [post_segment])


def mlm_seq(seq):
    seq_token, masked_sequence_token = [7],[7]
    seq_token += [tok_to_idx[token] for token in seq]
    
    masked_seq = random_replace(seq, n_mut) # 随机替换n_mut个元素为'*'
    masked_seq_token += [tok_to_idx[token] for token in masked_seq]
    
    return seq, masked_seq, torch.LongTensor(seq_token), torch.LongTensor(masked_seq_token)

def batch_mlm_seq(seq_list, continuous_replace = False):
    batch_seq = []
    batch_masked_seq = []
    batch_seq_token_list = []
    batch_masked_seq_token_list = []

    for i, seq in enumerate(seq_list):
        seq_token, masked_seq_token = [7], [7]
        seq_token += [tok_to_idx[token] for token in seq]

        masked_seq = random_replace(seq, continuous_replace)  # 随机替换n_mut个元素为'*'
        masked_seq_token += [tok_to_idx[token] for token in masked_seq]

        batch_seq.append(seq)
        batch_masked_seq.append(masked_seq)

        batch_seq_token_list.append(seq_token)
        batch_masked_seq_token_list.append(masked_seq_token)

    return batch_seq, batch_masked_seq, torch.LongTensor(batch_seq_token_list), torch.LongTensor(batch_masked_seq_token_list)

def recovered_mlm_tokens(masked_seqs, masked_toks, esm_logits, exclude_low_prob = False):
    # Only remain the AGCT logits
    esm_logits = esm_logits[:,:,3:7]
    # Get the predicted tokens using argmax
    predicted_toks = (esm_logits.argmax(dim=-1)+3).tolist()
    
    batch_size, seq_len, vocab_size = esm_logits.size()
    if exclude_low_prob: min_prob = 1 / vocab_size
    # Initialize an empty list to store the recovered sequences
    recovered_sequences, recovered_toks = [], []
    
    for i in range(batch_size):
        recovered_sequence_i, recovered_tok_i = [], []
        for j in range(seq_len):
            if masked_toks[i][j] == 8:
                st.write(i,j)
                ### Sample M recovery sequences using the logits
                recovery_probs = torch.softmax(esm_logits[i, j], dim=-1)
                recovery_probs[predicted_toks[i][j]-3] = 0  # Exclude the most probable token
                if exclude_low_prob: recovery_probs[recovery_probs < min_prob] = 0  # Exclude tokens with low probs < min_prob
                recovery_probs /= recovery_probs.sum()  # Normalize the probabilities
                
                ### 有放回抽样
                max_retries = 5
                retries = 0
                success = False

                while retries < max_retries and not success:
                    try:
                        recovery_indices = list(np.random.choice(vocab_size, size=n_mlm_recovery_sampling, p=recovery_probs.cpu().detach().numpy(), replace=False))
                        success = True  # 设置成功标志
                    except ValueError as e:
                        retries += 1
                        st.write(f"Attempt {retries} failed with error: {e}")
                        if retries >= max_retries:
                            st.write("Max retries reached. Skipping this iteration.")
                
                ### recovery to sequence
                if retries < max_retries:
                    for idx in [predicted_toks[i][j]] + [3+i for i in recovery_indices]:
                        recovery_seq = deepcopy(list(masked_seqs[i]))
                        recovery_tok = deepcopy(masked_toks[i])
                        
                        recovery_tok[j] = idx
                        recovery_seq[j-1] = idx_to_tok[idx]
                        
                        recovered_tok_i.append(recovery_tok)
                        recovered_sequence_i.append(''.join(recovery_seq))
                        
        recovered_sequences.extend(recovered_sequence_i)
        recovered_toks.extend(recovered_tok_i)
    return recovered_sequences, torch.LongTensor(torch.stack(recovered_toks))

def recovered_mlm_multi_tokens(masked_seqs, masked_toks, esm_logits, exclude_low_prob = False):
    # Only remain the AGCT logits
    esm_logits = esm_logits[:,:,3:7]
    # Get the predicted tokens using argmax
    predicted_toks = (esm_logits.argmax(dim=-1)+3).tolist()
    
    batch_size, seq_len, vocab_size = esm_logits.size()
    if exclude_low_prob: min_prob = 1 / vocab_size
    # Initialize an empty list to store the recovered sequences
    recovered_sequences, recovered_toks = [], []
    
    for i in range(batch_size):
        recovered_sequence_i, recovered_tok_i = [], []
        recovered_masked_num = 0
        for j in range(seq_len):
            if masked_toks[i][j] == 8:
                ### Sample M recovery sequences using the logits
                recovery_probs = torch.softmax(esm_logits[i, j], dim=-1)
                recovery_probs[predicted_toks[i][j]-3] = 0  # Exclude the most probable token
                if exclude_low_prob: recovery_probs[recovery_probs < min_prob] = 0  # Exclude tokens with low probs < min_prob
                recovery_probs /= recovery_probs.sum()  # Normalize the probabilities
                
                ### 有放回抽样
                max_retries = 5
                retries = 0
                success = False

                while retries < max_retries and not success:
                    try:
                        recovery_indices = list(np.random.choice(vocab_size, size=n_mlm_recovery_sampling, p=recovery_probs.cpu().detach().numpy(), replace=False))
                        success = True  # 设置成功标志
                    except ValueError as e:
                        retries += 1
                        st.write(f"Attempt {retries} failed with error: {e}")
                        if retries >= max_retries:
                            st.write("Max retries reached. Skipping this iteration.")
                
                ### recovery to sequence 
                        
                if recovered_masked_num == 0:
                    if retries < max_retries:
                        for idx in [predicted_toks[i][j]] + [3+i for i in recovery_indices]:
                            recovery_seq = deepcopy(list(masked_seqs[i]))
                            recovery_tok = deepcopy(masked_toks[i])

                            recovery_tok[j] = idx
                            recovery_seq[j-1] = idx_to_tok[idx]

                            recovered_tok_i.append(recovery_tok)
                            recovered_sequence_i.append(''.join(recovery_seq))

                elif recovered_masked_num > 0:
                    if retries < max_retries:
                        for idx in [predicted_toks[i][j]] + [3+i for i in recovery_indices]:
                            for recovery_seq, recovery_tok in zip(list(recovered_sequence_i), list(recovered_tok_i)): # 要在循环开始之前获取列表的副本来进行迭代。这样,在循环中即使我们修改了原始的列表,也不会影响迭代的行为。

                                recovery_seq_temp = list(recovery_seq)
                                recovery_tok[j] = idx
                                recovery_seq_temp[j-1] = idx_to_tok[idx]
                                
                                recovered_tok_i.append(recovery_tok)
                                recovered_sequence_i.append(''.join(recovery_seq_temp))
                
                recovered_masked_num += 1                
        recovered_indices = [i for i, s in enumerate(recovered_sequence_i) if '*' not in s]
        recovered_tok_i = [recovered_tok_i[i] for i in recovered_indices]
        recovered_sequence_i = [recovered_sequence_i[i] for i in recovered_indices]

        recovered_sequences.extend(recovered_sequence_i)
        recovered_toks.extend(recovered_tok_i)
        
        recovered_sequences, recovered_toks = remove_duplicates_double(recovered_sequences, recovered_toks)

    return recovered_sequences, torch.LongTensor(torch.stack(recovered_toks))

def mismatched_positions(s1, s2):
    # 这个函数假定两个字符串的长度相同。
    """Return the number of positions where two strings differ."""
    
    # The number of mismatches will be the sum of positions where characters are not the same
    return sum(1 for c1, c2 in zip(s1, s2) if c1 != c2)

def remove_duplicates_triple(filtered_mut_seqs, filtered_mut_probs, filtered_mut_logits):
    seen = {}
    unique_seqs = []
    unique_probs = []
    unique_logits = []
    
    for seq, prob, logit in zip(filtered_mut_seqs, filtered_mut_probs, filtered_mut_logits):
        if seq not in seen:
            unique_seqs.append(seq)
            unique_probs.append(prob)
            unique_logits.append(logit)
            seen[seq] = True

    return unique_seqs, unique_probs, unique_logits

def remove_duplicates_double(filtered_mut_seqs, filtered_mut_probs):
    seen = {}
    unique_seqs = []
    unique_probs = []
    
    for seq, prob in zip(filtered_mut_seqs, filtered_mut_probs):
        if seq not in seen:
            unique_seqs.append(seq)
            unique_probs.append(prob)
            seen[seq] = True

    return unique_seqs, unique_probs

def mutated_seq(wt_seq, wt_label):
    wt_seq = '!'+ wt_seq
    wt_tok = torch.LongTensor([[tok_to_idx[token] for token in wt_seq]]).to(device)
    wt_prob, wt_pred, _, wt_logit = model.predict(wt_tok)
    
    st.write(f'Wild Type: Length = {len(wt_seq)} \n{wt_seq}')
    st.write(f'Wild Type: Label = {wt_label}, Y_pred = {wt_pred.item()}, Y_prob = {wt_prob.item():.2%}')
    
#     st.write(n_mut, mlm_tok_num, n_designs_ep, n_sampling_designs_ep, n_mlm_recovery_sampling, mutate2stronger)
    # pbar = tqdm(total=n_mut)
    mutated_seqs = []
    i = 1
    # pbar = st.progress(i, text="mutated number of sequence")
    while i <= n_mut:
        if i == 1: seeds_ep = [wt_seq[1:]]
        seeds_next_ep, seeds_probs_next_ep, seeds_logits_next_ep = [], [], []
        for seed in seeds_ep:
            seed_seq, masked_seed_seq, seed_seq_token, masked_seed_seq_token = batch_mlm_seq([seed] * n_designs_ep, continuous_replace = True) ### mask seed with 1 site to "*"

            seed_prob, seed_pred, _, seed_logit = model.predict(seed_seq_token[0].unsqueeze_(0).to(device))
            _, _, seed_esm_logit, _ = model.predict(masked_seed_seq_token.to(device))
            mut_seqs, mut_toks = recovered_mlm_multi_tokens(masked_seed_seq, masked_seed_seq_token, seed_esm_logit)
            mut_probs, mut_preds, mut_esm_logits, mut_logits = model.predict(mut_toks.to(device))

            ### Filter mut_seqs that mut_prob < seed_prob and mut_prob < wild_prob
            filtered_mut_seqs = []
            filtered_mut_probs = []
            filtered_mut_logits = []
            if mut_by_prob:
                for z in range(len(mut_seqs)):
                    if mutate2stronger:
                        if mut_probs[z] >= seed_prob and mut_probs[z] >= wt_prob:
                            filtered_mut_seqs.append(mut_seqs[z])
                            filtered_mut_probs.append(mut_probs[z].cpu().detach().numpy())
                            filtered_mut_logits.append(mut_logits[z].cpu().detach().numpy())
                    else:
                        if mut_probs[z] < seed_prob and mut_probs[z] < wt_prob:
                            filtered_mut_seqs.append(mut_seqs[z])
                            filtered_mut_probs.append(mut_probs[z].cpu().detach().numpy())
                            filtered_mut_logits.append(mut_logits[z].cpu().detach().numpy())
            else:
                for z in range(len(mut_seqs)):
                    if mutate2stronger:
                        if mut_logits[z] >= seed_logit and mut_logits[z] >= wt_logit:
                            filtered_mut_seqs.append(mut_seqs[z])
                            filtered_mut_probs.append(mut_probs[z].cpu().detach().numpy())
                            filtered_mut_logits.append(mut_logits[z].cpu().detach().numpy())
                    else:
                        if mut_logits[z] < seed_logit and mut_logits[z] < wt_logit:
                            filtered_mut_seqs.append(mut_seqs[z])
                            filtered_mut_probs.append(mut_probs[z].cpu().detach().numpy())
                            filtered_mut_logits.append(mut_logits[z].cpu().detach().numpy())
                
            

            ### Save
            seeds_next_ep.extend(filtered_mut_seqs)
            seeds_probs_next_ep.extend(filtered_mut_probs)
            seeds_logits_next_ep.extend(filtered_mut_logits)
            seeds_next_ep, seeds_probs_next_ep, seeds_logits_next_ep = remove_duplicates_triple(seeds_next_ep, seeds_probs_next_ep, seeds_logits_next_ep)

        ### Sampling based on prob
        if len(seeds_next_ep) > n_sampling_designs_ep:
            seeds_probs_next_ep_norm = seeds_probs_next_ep / sum(seeds_probs_next_ep)  # Normalize the probabilities
            seeds_index_next_ep = np.random.choice(len(seeds_next_ep), n_sampling_designs_ep, p = seeds_probs_next_ep_norm, replace = False)
            
            seeds_next_ep = np.array(seeds_next_ep)[seeds_index_next_ep]
            seeds_probs_next_ep = np.array(seeds_probs_next_ep)[seeds_index_next_ep]
            seeds_logits_next_ep = np.array(seeds_logits_next_ep)[seeds_index_next_ep]
        seeds_mutated_num_next_ep = [mismatched_positions(wt_seq[1:], s) for s in seeds_next_ep]
        
        mutated_seqs.extend(list(zip(seeds_next_ep, seeds_logits_next_ep, seeds_probs_next_ep, seeds_mutated_num_next_ep)))

        seeds_ep = seeds_next_ep
        i += 1
        # pbar.update(1)
        # pbar.progress(i/n_mut, text="Mutating")
    # pbar.close()
    # st.success('Done', icon="✅")
    mutated_seqs.extend([(wt_seq[1:], wt_logit.item(), wt_prob.item(), 0)])
    mutated_seqs = sorted(mutated_seqs, key=lambda x: x[2], reverse=True)
    mutated_seqs = pd.DataFrame(mutated_seqs, columns = ['mutated_seq', 'predicted_logit', 'predicted_probability', 'mutated_num']).drop_duplicates('mutated_seq')
    return mutated_seqs

def read_raw(raw_input):
    ids = []
    sequences = []

    file = StringIO(raw_input)
    for record in SeqIO.parse(file, "fasta"):

        # 检查序列是否只包含A, G, C, T
        sequence = str(record.seq.back_transcribe()).upper()
        if not set(sequence).issubset(set("AGCT")):
            st.write(f"Record '{record.description}' was skipped for containing invalid characters. Only A, G, C, T(U) are allowed.")
            continue

        # 将符合条件的序列添加到列表中
        ids.append(record.id)
        sequences.append(sequence)
    
    return ids, sequences

def predict_raw(raw_input):
    model.eval()
    # st.write('====Parse Input====')
    ids, seqs = read_raw(raw_input)

    # st.write('====Predict====')
    res_pd = pd.DataFrame(columns = ['wildtype_id', 'mutated_seq', 'predicted_logit', 'predicted_probability', 'mutated_num'])
    for wt_seq, wt_id in zip(seqs, ids):
        try:
            res = mutated_seq(wt_seq, wt_id)
            res['wildtype_id'] = wt_id
            res_pd = pd.concat([res_pd,res], axis = 0)
        except:
            st.write('====Please Try Again this sequence: ', wt_id, wt_seq)
    
    # st.write(res_pd)
    return res_pd

global model, device
device = "cpu"
state_dict = torch.load('model.pt', map_location=torch.device(device))
new_state_dict = OrderedDict()

for k, v in state_dict.items():
    name = k.replace('module.','')
    new_state_dict[name] = v

model = CNN_linear().to(device)
model.load_state_dict(new_state_dict, strict = False)

# Run    
if st.button("Predict and Mutate"):
    if uploaded:
        result = predict_raw(uploaded.getvalue().decode())
    else:
        result = predict_raw(seq)
    
    result_file = result.to_csv(index=False)
    st.download_button("Download", result_file, file_name=output_filename+".csv")
    st.dataframe(result)