Upload vocab.py

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