Upload 5 files

bertmodel.py

@@ -0,0 +1,171 @@
+import torch.nn as nn
+import copy, math
+import torch
+import numpy as np
+import torch.nn.functional as F
+
+class Bert(nn.Module):
+    
+    def __init__(self, encoder, src_embed):
+        super(Bert, self).__init__()
+
+        self.encoder = encoder        
+        self.src_embed = src_embed  
+    
+    def forward(self, src, src_mask):
+    
+        return self.encoder(self.src_embed(src), src_mask)
+    
+    
+class Encoder(nn.Module):
+    def __init__(self, layer, N):
+        super(Encoder, self).__init__()
+        self.layers = clones(layer, N)
+        self.norm = LayerNorm(layer.size)
+        
+    def forward(self, x, mask):
+        for layer in self.layers:
+            x = layer(x, mask)
+        return self.norm(x)
+    
+class LayerNorm(nn.Module):
+    def __init__(self, features, eps=1e-6):
+        super(LayerNorm, self).__init__()
+        self.a_2 = nn.Parameter(torch.ones(features))
+        self.b_2 = nn.Parameter(torch.zeros(features))
+        self.eps = eps
+
+    def forward(self, x):
+        mean = x.mean(-1, keepdim=True)
+        std = x.std(-1, keepdim=True)
+        return self.a_2 * (x - mean) / (std + self.eps) + self.b_2
+
+class SublayerConnection(nn.Module):
+    def __init__(self, size, dropout):
+        super(SublayerConnection, self).__init__()
+        self.norm = LayerNorm(size)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x, sublayer):
+        return x + self.dropout(sublayer(self.norm(x)))
+    
+class EncoderLayer(nn.Module):
+    def __init__(self, size, self_attn, feed_forward, dropout):
+        super(EncoderLayer, self).__init__()
+        self.self_attn = self_attn
+        self.feed_forward = feed_forward
+        self.sublayer = clones(SublayerConnection(size, dropout), 2)
+        self.size = size
+
+    def forward(self, x, mask):
+        x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, mask))
+        return self.sublayer[1](x, self.feed_forward)
+    
+class PositionwiseFeedForward(nn.Module):
+    def __init__(self, d_model, d_ff, dropout=0.1):
+        super(PositionwiseFeedForward, self).__init__()
+        self.w_1 = nn.Linear(d_model, d_ff)
+        self.w_2 = nn.Linear(d_ff, d_model)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x):
+        return self.w_2(self.dropout(F.relu(self.w_1(x))))
+    
+def make_bert(src_vocab, N=6, d_model=512, d_ff=2048, h=8, dropout=0.1):
+    c = copy.deepcopy
+    attn = MultiHeadedAttention(h, d_model)
+    ff = PositionwiseFeedForward(d_model, d_ff, dropout)
+    position = PositionalEncoding(d_model, dropout)
+    model = Bert(
+        Encoder(EncoderLayer(d_model, c(attn), c(ff), dropout), N),
+        
+        nn.Sequential(Embeddings(d_model, src_vocab), c(position)),
+    )
+    
+    for p in model.parameters():
+        if p.dim() > 1:
+            nn.init.xavier_uniform_(p)
+    return model
+
+def make_bert_without_emb(d_model=128, N=2, d_ff=512, h=8, dropout=0.1):
+    c = copy.deepcopy
+    attn = MultiHeadedAttention(h, d_model)
+    ff = PositionwiseFeedForward(d_model, d_ff, dropout)
+    trainable_encoder = Encoder(EncoderLayer(d_model, c(attn), c(ff), dropout), N)
+
+    return trainable_encoder
+
+
+
+def clones(module, N):
+    return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])
+
+def subsequent_mask(size):
+    attn_shape = (1, size, size)
+    subsequent_mask = np.triu(np.ones(attn_shape), k=1).astype('uint8')
+    return torch.from_numpy(subsequent_mask) == 0
+
+def attention(query, key, value, mask=None, dropout=None):
+    d_k = query.size(-1)
+    scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)
+    if mask is not None:
+        mask = mask.unsqueeze(-2)
+        scores = scores.masked_fill(mask == 0, -1e9)
+    p_attn = F.softmax(scores, dim = -1)
+    if dropout is not None:
+        p_attn = dropout(p_attn)
+    return torch.matmul(p_attn, value), p_attn
+
+class MultiHeadedAttention(nn.Module):
+    def __init__(self, h, d_model, dropout=0.1):
+        super(MultiHeadedAttention, self).__init__()
+        assert d_model % h == 0
+        self.d_k = d_model // h
+        self.h = h
+        self.linears = clones(nn.Linear(d_model, d_model), 4)
+        self.attn = None
+        self.dropout = nn.Dropout(p=dropout)
+        
+    def forward(self, query, key, value, mask=None):
+        if mask is not None:
+            mask = mask.unsqueeze(1)
+        nbatches = query.size(0)
+          
+        query, key, value = \
+            [l(x).view(nbatches, -1, self.h, self.d_k).transpose(1, 2)
+             for l, x in zip(self.linears, (query, key, value))]
+        
+        x, self.attn = attention(query, key, value, mask=mask, 
+                                 dropout=self.dropout)
+        
+        x = x.transpose(1, 2).contiguous() \
+             .view(nbatches, -1, self.h * self.d_k)
+        return self.linears[-1](x)
+    
+class Embeddings(nn.Module):
+    def __init__(self, d_model, vocab):
+        super(Embeddings, self).__init__()
+        self.lut = nn.Embedding(vocab, d_model)
+        self.d_model = d_model
+
+    def forward(self, x):
+        return self.lut(x) * math.sqrt(self.d_model)
+    
+class PositionalEncoding(nn.Module):
+    def __init__(self, d_model, dropout, max_len=5000):
+        super(PositionalEncoding, self).__init__()
+        self.dropout = nn.Dropout(p=dropout)
+        
+        pe = torch.zeros(max_len, d_model)
+        position = torch.arange(0, max_len).unsqueeze(1)
+        div_term = torch.exp(torch.arange(0, d_model, 2) *
+                             -(math.log(10000.0) / d_model))
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        pe = pe.unsqueeze(0)
+        self.register_buffer('pe', pe)
+        
+    def forward(self, x):
+        x = x + self.pe[:, :x.size(1)].clone().detach()
+        return self.dropout(x)
+    

dataset_mlm.py

@@ -0,0 +1,151 @@
+import pandas as pd
+from copy import deepcopy
+
+import torch
+from torch.utils.data import TensorDataset, DataLoader
+from sklearn.model_selection import train_test_split
+
+from vocab import PepVocab
+from utils import mask, create_vocab
+
+addtition_tokens = ['<K16>', '<α1β1γδ>', '<Ca22>', '<AChBP>', '<K13>', '<α1BAR>', '<α1β1ε>', '<α1AAR>', '<GluN3A>', '<α4β2>',
+                     '<GluN2B>', '<α75HT3>', '<Na14>', '<α7>', '<GluN2C>', '<NET>', '<NavBh>', '<α6β3β4>', '<Na11>', '<Ca13>', 
+                     '<Ca12>', '<Na16>', '<α6α3β2>', '<GluN2A>', '<GluN2D>', '<K17>', '<α1β1δε>', '<GABA>', '<α9>', '<K12>', 
+                     '<Kshaker>', '<α3β4>', '<Na18>', '<α3β2>', '<α6α3β2β3>', '<α1β1δ>', '<α6α3β4β3>', '<α2β2>','<α6β4>', '<α2β4>',
+                     '<Na13>', '<Na12>', '<Na15>', '<α4β4>', '<α7α6β2>', '<α1β1γ>', '<NaTTXR>', '<K11>', '<Ca23>', 
+                     '<α9α10>','<α6α3β4>', '<NaTTXS>', '<Na17>','<high>','<low>']
+
+def add_tokens_to_vocab(vocab_mlm: PepVocab):
+    vocab_mlm.add_special_token(addtition_tokens)
+    return vocab_mlm
+
+def split_seq(seq, vocab, get_seq=False):
+    '''
+    note: the function is suitable for the sequences with the format of "label|label|sequence|msa1|msa2|msa3"
+    '''
+    start = '[CLS]'
+    end = '[SEP]'
+    pad = '[PAD]'
+    cls_label = seq.split('|')[0]
+    act_label = seq.split('|')[1]
+
+    if get_seq == True:
+        add = lambda x: [start] + [cls_label] + [act_label] + x + [end]
+        pep_seq = seq.split('|')[2]
+        # return [start] + [cls_label] + [act_label] + vocab.split_seq(pep_seq) + [end]
+        return add(vocab.split_seq(pep_seq))
+    
+    else:
+        add = lambda x: [start] + [pad] + [pad] + x + [end]
+        msa1_seq = seq.split('|')[3]
+        msa2_seq = seq.split('|')[4]
+        msa3_seq = seq.split('|')[5]
+
+        # return [vocab.split_seq(msa1_seq)]  + [vocab.split_seq(msa2_seq)]  + [vocab.split_seq(msa3_seq)]
+        return [add(vocab.split_seq(msa1_seq))]  + [add(vocab.split_seq(msa2_seq))]  + [add(vocab.split_seq(msa3_seq))]
+
+def get_paded_token_idx(vocab_mlm):
+    cono_path = 'conoData_C5.csv'
+    seq = pd.read_csv(cono_path)['Sequences']
+    
+    splited_seq = list(seq.apply(split_seq, args=(vocab_mlm,True, )))
+    splited_msa = list(seq.apply(split_seq, args=(vocab_mlm, False, )))
+    
+    vocab_mlm.set_get_attn(is_get=True)
+    padded_seq = vocab_mlm.truncate_pad(splited_seq, num_steps=54, padding_token='[PAD]')
+    attn_idx = vocab_mlm.get_attention_mask_mat()
+
+    vocab_mlm.set_get_attn(is_get=False)
+    padded_msa = vocab_mlm.truncate_pad(splited_msa, num_steps=54, padding_token='[PAD]')
+    
+    idx_seq = vocab_mlm.__getitem__(padded_seq) # [b, 54]  start, cls_label, act_label, sequence, end
+    
+    idx_msa = vocab_mlm.__getitem__(padded_msa) # [b, 3, 50]
+
+    return padded_seq, idx_seq, idx_msa, attn_idx
+
+def get_paded_token_idx_gen(vocab_mlm, seq):
+    
+    splited_seq = split_seq(seq[0], vocab_mlm, True)
+    splited_msa = split_seq(seq[0], vocab_mlm, False)
+    
+    vocab_mlm.set_get_attn(is_get=True)
+    padded_seq = vocab_mlm.truncate_pad(splited_seq, num_steps=54, padding_token='[PAD]')
+    attn_idx = vocab_mlm.get_attention_mask_mat()
+
+    vocab_mlm.set_get_attn(is_get=False)
+    padded_msa = vocab_mlm.truncate_pad(splited_msa, num_steps=54, padding_token='[PAD]')
+    
+    idx_seq = vocab_mlm.__getitem__(padded_seq) # [b, 54]  start, cls_label, act_label, sequence, end
+    
+    idx_msa = vocab_mlm.__getitem__(padded_msa) # [b, 3, 50]
+
+    return padded_seq, idx_seq, idx_msa, attn_idx
+
+
+def get_paded_token_idx_gen(vocab_mlm, seq, new_seq):
+    if new_seq == None:
+        splited_seq = split_seq(seq[0], vocab_mlm, True)
+        splited_msa = split_seq(seq[0], vocab_mlm, False)
+        
+        vocab_mlm.set_get_attn(is_get=True)
+        padded_seq = vocab_mlm.truncate_pad(splited_seq, num_steps=54, padding_token='[PAD]')
+        attn_idx = vocab_mlm.get_attention_mask_mat()
+        vocab_mlm.set_get_attn(is_get=False)
+
+        padded_msa = vocab_mlm.truncate_pad(splited_msa, num_steps=54, padding_token='[PAD]')
+        
+        idx_seq = vocab_mlm.__getitem__(padded_seq)  # [b, 54]  start, cls_label, act_label, sequence, end
+        idx_msa = vocab_mlm.__getitem__(padded_msa)  # [b, 3, 50]
+    else:
+        splited_seq = split_seq(seq[0], vocab_mlm, True)
+        splited_msa = split_seq(seq[0], vocab_mlm, False)
+        vocab_mlm.set_get_attn(is_get=True)
+        padded_seq = vocab_mlm.truncate_pad(splited_seq, num_steps=54, padding_token='[PAD]')
+        attn_idx = vocab_mlm.get_attention_mask_mat()
+        vocab_mlm.set_get_attn(is_get=False)
+        padded_msa = vocab_mlm.truncate_pad(splited_msa, num_steps=54, padding_token='[PAD]')
+        idx_msa = vocab_mlm.__getitem__(padded_msa)  # [b, 3, 50]
+
+        idx_seq = vocab_mlm.__getitem__(new_seq)
+    return padded_seq, idx_seq, idx_msa, attn_idx
+
+
+
+def make_mask(seq_ser, start, end, time, vocab_mlm, labels, idx_msa, attn_idx):
+    seq_ser = pd.Series(seq_ser)
+    masked_seq = seq_ser.apply(mask, args=(start, end, time))
+    masked_idx = vocab_mlm.__getitem__(list(masked_seq))
+    masked_idx = torch.tensor(masked_idx)
+    device = torch.device('cuda:1')
+    data_arrays = (masked_idx.to(device), labels.to(device), idx_msa.to(device), attn_idx.to(device)) 
+    dataset = TensorDataset(*data_arrays)
+    train_dataset, test_dataset = train_test_split(dataset, test_size=0.1, random_state=42, shuffle=True)
+    train_loader = DataLoader(train_dataset, batch_size=128, shuffle=True)
+    test_loader = DataLoader(test_dataset, batch_size=128, shuffle=True)
+    
+    return train_loader, test_loader
+
+if __name__ == '__main__':
+    # from add_args import parse_args
+    import numpy as np
+    # args = parse_args()
+
+    vocab_mlm = create_vocab()
+    vocab_mlm = add_tokens_to_vocab(vocab_mlm)
+    padded_seq, idx_seq, idx_msa, attn_idx = get_paded_token_idx(vocab_mlm)
+    labels = torch.tensor(idx_seq)
+    idx_msa = torch.tensor(idx_msa)
+    attn_idx = torch.tensor(attn_idx)
+
+    # time_step = args.mask_time_step
+    for t in np.arange(1, 50):
+        padded_seq_copy = deepcopy(padded_seq)        
+        train_loader, test_loader = make_mask(padded_seq_copy, start=0, end=49, time=t, 
+                                              vocab_mlm=vocab_mlm, labels=labels, idx_msa=idx_msa, attn_idx=attn_idx)
+        for i, (masked_idx, label, msa, attn) in enumerate(train_loader):
+            print(f"the {i}th batch is that masked_idx is {masked_idx.shape}, labels is {label.shape}, idx_msa is {msa.shape}")
+        print(f"the {t}th time step is done")
+        
+        
+    

model.py

@@ -0,0 +1,171 @@
+import torch.nn as nn
+import copy, math
+import torch
+import numpy as np
+import torch.nn.functional as F
+from transformers import AutoModelForMaskedLM, AutoConfig
+
+from bertmodel import make_bert, make_bert_without_emb
+from utils import ContraLoss
+    
+def load_pretrained_model():
+    model_checkpoint = "Rostlab/prot_bert"
+    config = AutoConfig.from_pretrained(model_checkpoint)
+    model = AutoModelForMaskedLM.from_config(config)
+    
+    return model
+
+class ConoEncoder(nn.Module):
+    def __init__(self, encoder):
+        super(ConoEncoder, self).__init__()
+        
+        self.encoder = encoder
+        self.trainable_encoder = make_bert_without_emb()
+
+        
+        for param in self.encoder.parameters():
+            param.requires_grad = False
+        
+        
+    def forward(self, x, mask):  # x:(128,54)  mask:(128,54)
+        feat = self.encoder(x, attention_mask=mask)  # (128,54,128)
+        feat = list(feat.values())[0] # (128,54,128)
+        
+        feat = self.trainable_encoder(feat, mask) # (128,54,128)
+
+        return feat
+
+class MSABlock(nn.Module):
+    def __init__(self, in_dim, out_dim, vocab_size):
+        super(MSABlock, self).__init__()
+        self.embedding = nn.Embedding(vocab_size, in_dim)
+        self.mlp = nn.Sequential(
+            nn.Linear(in_dim, out_dim),
+            nn.LeakyReLU(),
+            nn.Linear(out_dim, out_dim)
+        )
+        self.init()
+    
+    def init(self):
+        for layer in self.mlp.children():
+            if isinstance(layer, nn.Linear):
+                nn.init.xavier_uniform_(layer.weight)
+        # nn.init.xavier_uniform_(self.embedding.weight)
+
+    def forward(self, x):  # x: (128,3,54)
+        x = self.embedding(x) # x: (128,3,54,128)
+        x = self.mlp(x) # x: (128,3,54,128)
+        return x
+
+class ConoModel(nn.Module):
+    def __init__(self, encoder, msa_block, decoder):
+        super(ConoModel, self).__init__()
+        self.encoder = encoder
+        self.msa_block = msa_block
+        self.feature_combine = nn.Conv2d(in_channels=4, out_channels=1, kernel_size=1)
+        self.decoder = decoder
+
+    def forward(self, input_ids, msa, attn_idx=None):
+        encoder_output = self.encoder.forward(input_ids, attn_idx) # (128,54,128)
+        msa_output = self.msa_block(msa) # (128,3,54,128)
+        # msa_output = torch.mean(msa_output, dim=1)
+        encoder_output = encoder_output.view(input_ids.shape[0], 54, -1).unsqueeze(1) # (128,1,54,128)
+        
+        output = torch.cat([encoder_output*5, msa_output], dim=1) # (128,4,54,128)
+        output = self.feature_combine(output) # (128,1,54,128)
+        output = output.squeeze(1) # (128,54,128)
+        logits = self.decoder(output) # (128,54,85)
+        
+        return logits
+
+class ContraModel(nn.Module):
+    def __init__(self, cono_encoder):
+        super(ContraModel, self).__init__()
+        
+        self.contra_loss = ContraLoss()
+
+        self.encoder1 = cono_encoder
+        self.encoder2 = make_bert(404, 6, 128)
+
+        # contrastive decoder
+        self.lstm = nn.LSTM(16, 16, batch_first=True)
+        self.contra_decoder = nn.Sequential(
+            nn.Linear(128, 64),
+            nn.LeakyReLU(),
+            nn.Linear(64, 32),
+            nn.LeakyReLU(),
+            nn.Linear(32, 16),
+            nn.LeakyReLU(),
+            nn.Dropout(0.1),
+        )
+        
+        # classifier
+        self.pre_classifer = nn.LSTM(128, 64, batch_first=True)
+        self.classifer = nn.Sequential(
+            nn.Linear(128, 32),
+            nn.LeakyReLU(),
+            nn.Linear(32, 6),
+            nn.Softmax(dim=-1)
+        )
+
+        self.init()
+
+    def init(self):
+        
+        for layer in self.contra_decoder.children():
+            if isinstance(layer, nn.Linear):
+                nn.init.xavier_uniform_(layer.weight)
+        for layer in self.classifer.children():
+            if isinstance(layer, nn.Linear):
+                nn.init.xavier_uniform_(layer.weight)
+        for layer in self.pre_classifer.children():
+            if isinstance(layer, nn.Linear):
+                nn.init.xavier_uniform_(layer.weight)
+        for layer in self.lstm.children():
+            if isinstance(layer, nn.Linear):
+                nn.init.xavier_uniform_(layer.weight)
+
+    def compute_class_loss(self, feat1, feat2, labels):
+        _, cls_feat1= self.pre_classifer(feat1)
+        _, cls_feat2 = self.pre_classifer(feat2)
+        cls_feat1 = torch.cat([cls_feat1[0], cls_feat1[1]], dim=-1).squeeze(0)
+        cls_feat2 = torch.cat([cls_feat2[0], cls_feat2[1]], dim=-1).squeeze(0)
+
+        cls1_dis = self.classifer(cls_feat1)
+        cls2_dis = self.classifer(cls_feat2)
+        cls1_loss = F.cross_entropy(cls1_dis, labels.to('cuda:0'))
+        cls2_loss = F.cross_entropy(cls2_dis, labels.to('cuda:0'))
+        
+        return cls1_loss, cls2_loss
+
+    def compute_contrastive_loss(self, feat1, feat2):
+    
+        contra_feat1 = self.contra_decoder(feat1)
+        contra_feat2 = self.contra_decoder(feat2)
+        
+        _, feat1 = self.lstm(contra_feat1)
+        _, feat2 = self.lstm(contra_feat2)
+        feat1 = torch.cat([feat1[0], feat1[1]], dim=-1).squeeze(0)
+        feat2 = torch.cat([feat2[0], feat2[1]], dim=-1).squeeze(0)
+
+        ctr_loss = self.contra_loss(feat1, feat2)
+    
+        return ctr_loss
+    
+    def forward(self, x1, x2, labels=None):
+        loss = dict()
+
+        idx1, attn1 = x1
+        idx2, attn2 = x2
+        feat1 = self.encoder1(idx1.to('cuda:0'), attn1.to('cuda:0'))
+        feat2 = self.encoder2(idx2.to('cuda:0'), attn2.to('cuda:0'))
+        
+        cls1_loss, cls2_loss = self.compute_class_loss(feat1, feat2, labels)
+
+        ctr_loss = self.compute_contrastive_loss(feat1, feat2)
+
+        loss['cls1_loss'] = cls1_loss
+        loss['cls2_loss'] = cls2_loss
+        loss['ctr_loss'] = ctr_loss
+
+        return loss

utils.py

@@ -0,0 +1,132 @@
+import torch.nn as nn
+import copy, math
+import torch
+import numpy as np
+import torch.nn.functional as F
+
+from vocab import PepVocab
+
+def create_vocab():
+    vocab_mlm = PepVocab()
+    vocab_mlm.vocab_from_txt('vocab.txt')
+    # vocab_mlm.token_to_idx['-'] = 23
+    return vocab_mlm
+
+def show_parameters(model: nn.Module, show_all=False, show_trainable=True):
+
+    mlp_pa = {name:param.requires_grad for name, param in model.named_parameters()}
+    
+    if show_all:
+        print('All parameters:')
+        print(mlp_pa)
+
+    if show_trainable:
+        print('Trainable parameters:')
+        print(list(filter(lambda x: x[1], list(mlp_pa.items()))))
+
+class ContraLoss(nn.Module):
+    def __init__(self, *args, **kwargs) -> None:
+        super(ContraLoss, self).__init__(*args, **kwargs)
+        
+        self.temp = 0.07
+
+    def contrastive_loss(self, proj1, proj2):
+        proj1 = F.normalize(proj1, dim=1)
+        proj2 = F.normalize(proj2, dim=1)
+        dot = torch.matmul(proj1, proj2.T) / self.temp
+        dot_max, _ = torch.max(dot, dim=1, keepdim=True)
+        dot = dot - dot_max.detach()
+
+        exp_dot = torch.exp(dot)
+        log_prob = torch.diag(dot, 0) - torch.log(exp_dot.sum(1))
+        cont_loss = -log_prob.mean()
+        return cont_loss
+    
+    def forward(self, x, y, label=None):
+        return self.contrastive_loss(x, y)
+
+
+import numpy as np
+from tqdm import tqdm
+import torch
+import torch.nn as nn
+import random
+from transformers import set_seed
+
+def show_parameters(model: nn.Module, show_all=False, show_trainable=True):
+
+    mlp_pa = {name:param.requires_grad for name, param in model.named_parameters()}
+    
+    if show_all:
+        print('All parameters:')
+        print(mlp_pa)
+
+    if show_trainable:
+        print('Trainable parameters:')
+        print(list(filter(lambda x: x[1], list(mlp_pa.items()))))
+
+def extract_args(text):
+    str_list = []
+    substr = ""
+    for s in text:
+        if s in ('(', ')', '=', ',', ' ', '\n', "'"):
+            if substr != '':
+                str_list.append(substr)
+                substr = ''
+        else:
+            substr += s
+
+def eval_one_epoch(loader, cono_encoder):
+    cono_encoder.eval()
+    batch_loss = []
+    for i, data in enumerate(tqdm(loader)):
+        
+        loss = cono_encoder.contra_forward(data)
+        batch_loss.append(loss.item())
+        print(f'[INFO] Test batch {i} loss: {loss.item()}')
+
+    total_loss = np.mean(batch_loss)    
+    print(f'[INFO] Total loss: {total_loss}')    
+    return total_loss
+
+def setup_seed(seed):
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    np.random.seed(seed)
+    random.seed(seed)
+    torch.backends.cudnn.deterministic = True
+    set_seed(seed)
+
+class CrossEntropyLossWithMask(torch.nn.Module):
+    def __init__(self, weight=None):
+        super(CrossEntropyLossWithMask, self).__init__()
+        self.criterion = nn.CrossEntropyLoss(reduction='none')
+
+    def forward(self, y_pred, y_true, mask):
+        (pos_mask, label_mask, seq_mask) = mask
+        loss = self.criterion(y_pred, y_true) # (6912)
+        
+        pos_loss = (loss * pos_mask).sum() / torch.sum(pos_mask)
+        label_loss = (loss * label_mask).sum() / torch.sum(label_mask)
+        seq_loss = (loss * seq_mask).sum() / torch.sum(seq_mask)
+        
+        loss = pos_loss + label_loss/2 + seq_loss/3
+
+        return loss
+
+
+def mask(x, start, end, time):
+    ske_pos = np.where(np.array(x)=='C')[0] - start
+    lables_pos = np.array([1, 2]) - start
+    ske_pos = list(filter(lambda x: end-start >= x >= 0, ske_pos))
+    lables_pos = list(filter(lambda x: x >= 0, lables_pos))
+    weight = np.ones(end - start+1)
+    rand = np.random.rand()
+    if rand < 0.5:
+        weight[lables_pos] = 100000
+    else:
+        weight[lables_pos] = 1
+    mask_pos = np.random.choice(range(start, end+1), time, p=weight/np.sum(weight), replace=False)
+    for idx in mask_pos:
+        x[idx]  = '[MASK]'
+    return x

vocab.py

@@ -0,0 +1,193 @@
+import re
+import pandas as pd
+
+class PepVocab:
+    def __init__(self):
+        self.token_to_idx = { 
+            '<MASK>': -1, '<PAD>': 0, 'A': 1, 'C': 2, 'E': 3, 'D': 4, 'F': 5, 'I': 6, 'H': 7,
+            'K': 8, 'M': 9, 'L': 10, 'N': 11, 'Q': 12, 'P': 13, 'S': 14,
+            'R': 15, 'T': 16, 'W': 17, 'V': 18, 'Y': 19, 'G': 20, 'O': 21, 'U': 22, 'Z': 23, 'X': 24}
+        self.idx_to_token = { 
+            -1: '<MASK>', 0: '<PAD>', 1: 'A', 2: 'C', 3: 'E', 4: 'D', 5: 'F', 6: 'I', 7: 'H',
+            8: 'K', 9: 'M', 10: 'L', 11: 'N', 12: 'Q', 13: 'P', 14: 'S',
+            15: 'R', 16: 'T', 17: 'W', 18: 'V', 19: 'Y', 20: 'G', 21: 'O', 22: 'U', 23: 'Z', 24: 'X'}
+        
+        self.get_attention_mask = False
+        self.attention_mask = []
+        
+    def set_get_attn(self, is_get: bool):
+        self.get_attention_mask = is_get
+
+    def __len__(self):
+        return len(self.idx_to_token)
+
+    def __getitem__(self, tokens):
+        '''
+        note: input should a splited sequence
+
+        Args:
+            tokens: a token or token list of splited
+        '''
+        if not isinstance(tokens, (list, tuple)):
+            # return self.token_to_idx.get(tokens)
+            return self.token_to_idx[tokens]
+        return [self.__getitem__(token) for token in tokens]
+    
+    def vocab_from_txt(self, path):
+        '''
+        note: this function use for constructing vocab mapping
+        but it is only suitable for special txt format
+        it support one column txt file, which column name is 0
+        '''
+        token_to_idx = {}
+        idx_to_token = {}
+        chr_idx = pd.read_csv(path, header=None, sep='\t')
+        if chr_idx.shape[1] == 1:
+            for idx, token in enumerate(chr_idx[0]):
+                token_to_idx[token] = idx
+                idx_to_token[idx] = token
+        self.token_to_idx = token_to_idx
+        self.idx_to_token = idx_to_token
+        
+    def to_tokens(self, indices):
+        '''
+        note: input should a integer list
+        '''
+        if hasattr(indices, '__len__') and len(indices) > 1:
+            return [self.idx_to_token[int(index)] for index in indices]
+        return self.idx_to_token[indices]
+    
+    def add_special_token(self, token: str|list|tuple) -> None:
+        if not isinstance(token, (list, tuple)):
+            if token in self.token_to_idx:
+                raise ValueError(f"token {token} already in the vocab")
+            self.idx_to_token[len(self.idx_to_token)] = token
+            self.token_to_idx[token] = len(self.token_to_idx)
+        else:
+            [self.add_special_token(t) for t in token]
+        
+    def split_seq(self, seq: str|list|tuple) -> list:
+        if not isinstance(seq, (list, tuple)):
+            return re.findall(r"<[a-zA-Z0-9]+>|[a-zA-Z-]", seq)
+        return [self.split_seq(s) for s in seq] # a list of list
+    
+    def truncate_pad(self, line, num_steps, padding_token='<PAD>') -> list:
+
+        if not isinstance(line[0], list):
+            if len(line) > num_steps:
+                if self.get_attention_mask:
+                    self.attention_mask.append([1]*num_steps)
+                return line[:num_steps]
+            if self.get_attention_mask:
+                self.attention_mask.append([1] * len(line) + [0] * (num_steps - len(line)))
+            return line + [padding_token] * (num_steps - len(line))
+        else:
+            return [self.truncate_pad(l, num_steps, padding_token) for l in line]   # a list of list
+    
+    def get_attention_mask_mat(self):
+        attention_mask = self.attention_mask
+        self.attention_mask = []
+        return attention_mask
+
+    def seq_to_idx(self, seq: str|list|tuple, num_steps: int, padding_token='<PAD>') -> list:
+        '''
+        note: ensure to execut this function after add_special_token
+        '''
+
+        splited_seq = self.split_seq(seq)
+        # **********************
+        # after split, we need to mask sequence
+        # note: 
+        # 1. mask tokens by probability
+        # 2. return a list or list of list
+        # **********************
+        padded_seq = self.truncate_pad(splited_seq, num_steps, padding_token)
+
+        return self.__getitem__(padded_seq)
+
+
+
+class MutilVocab:
+    def __init__(self, data, AA_tok_len=2):
+        """
+        Args:
+            data (_type_):
+            AA_tok_len (int, optional): Defaults to 1.
+            start_token (bool, optional): True is required for encoder-based model.
+        """
+        ## Load train dataset
+        self.x_data = data
+        self.tok_AA_len = AA_tok_len
+        self.default_AA = list("RHKDESTNQCGPAVILMFYW")
+        # AAs which are not included in default_AA
+        self.tokens = self._token_gen(self.tok_AA_len)
+
+        self.token_to_idx = {k: i + 4 for i, k in enumerate(self.tokens)}
+        self.token_to_idx["[PAD]"] = 0  ## idx as 0 is PAD
+        self.token_to_idx["[CLS]"] = 1  ## idx as 1 is CLS
+        self.token_to_idx["[SEP]"] = 2  ## idx as 2 is SEP
+        self.token_to_idx["[MASK]"] = 3  ## idx as 3 is MASK
+        
+    def split_seq(self):
+        self.X = [self._seq_to_tok(seq) for seq in self.x_data]
+        return self.X
+    
+    def tok_idx(self, seqs):
+        '''
+        note: ensure to execut this function before truancate_pad
+        '''
+
+        seqs_idx = []
+        for seq in seqs:
+            seq_idx = []
+            for s in seq:
+                seq_idx.append(self.token_to_idx[s])
+            seqs_idx.append(seq_idx)
+
+        return seqs_idx
+
+
+
+    def _token_gen(self, tok_AA_len: int, st: str = "", curr_depth: int = 0):
+        """Generate tokens based on default amino acid residues
+            and also includes "X" as arbitrary residues.
+            Length of AAs in each token should be provided by "tok_AA_len"
+
+        Args:
+            tok_AA_len (int): Length of token
+            st (str, optional): Defaults to ''.
+            curr_depth (int, optional): Defaults to 0.
+
+        Returns:
+            List: List of tokens
+        """
+        curr_depth += 1
+        if curr_depth <= tok_AA_len:
+            l = [
+                st + t
+                for s in self.default_AA
+                for t in self._token_gen(tok_AA_len, s, curr_depth)
+            ]
+            return l
+        else:
+            return [st]
+
+    def _seq_to_tok(self, seq: str):
+        """Convert each token to index
+
+        Args:
+            seq (str): AA sequence
+
+        Returns:
+            list: A list of indexes
+        """
+
+        seq_idx = []
+        
+        seq_idx += ["[CLS]"]
+
+        for i in range(len(seq) - self.tok_AA_len + 1):
+            curr_token = seq[i : i + self.tok_AA_len]
+            seq_idx.append(curr_token)
+        seq_idx += ['[SEP]']
+        return seq_idx