dataset_mlm.py

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")