File size: 4,563 Bytes
2618f79 |
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 |
import itertools
from transformers import PreTrainedTokenizer
import json
import os
class KmerTokenizer(PreTrainedTokenizer):
def __init__(self, vocab_file=None, kmerlen=6, overlapping=True, maxlen=400, **kwargs):
self.kmerlen = kmerlen
self.overlapping = overlapping
self.maxlen = maxlen
# Initialize vocabulary
self.VOCAB = [''.join(i) for i in itertools.product(*(['ATCG'] * int(self.kmerlen)))]
self.VOCAB_SIZE = len(self.VOCAB) + 5
self.tokendict = dict(zip(self.VOCAB, range(5, self.VOCAB_SIZE)))
self.tokendict['[UNK]'] = 0
self.tokendict['[SEP]'] = 1
self.tokendict['[CLS]'] = 2
self.tokendict['[MASK]'] = 3
self.tokendict['[PAD]'] = 4
super().__init__(**kwargs)
def _tokenize(self, text):
tokens = []
stoprange = len(text) - (self.kmerlen - 1)
if self.overlapping:
for k in range(0, stoprange):
kmer = text[k:k + self.kmerlen]
if set(kmer).issubset('ATCG'):
tokens.append(kmer)
else:
for k in range(0, stoprange, self.kmerlen):
kmer = text[k:k + self.kmerlen]
if set(kmer).issubset('ATCG'):
tokens.append(kmer)
return tokens
def _convert_token_to_id(self, token):
return self.tokendict.get(token, self.tokendict['[UNK]'])
def _convert_id_to_token(self, index):
inv_tokendict = {v: k for k, v in self.tokendict.items()}
return inv_tokendict.get(index, '[UNK]')
def convert_tokens_to_string(self, tokens):
return ' '.join(tokens)
def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None):
if token_ids_1 is None:
return [self.tokendict['[CLS]']] + token_ids_0 + [self.tokendict['[SEP]']]
return [self.tokendict['[CLS]']] + token_ids_0 + [self.tokendict['[SEP]']] + token_ids_1 + [self.tokendict['[SEP]']]
def get_vocab(self):
return self.tokendict
def kmer_tokenize(self, seq_list):
seq_ind_list = []
for seq in seq_list:
tokens = self._tokenize(seq)
token_ids = [self._convert_token_to_id(token) for token in tokens]
if len(token_ids) < self.maxlen:
token_ids.extend([self.tokendict['[PAD]']] * (self.maxlen - len(token_ids)))
else:
token_ids = token_ids[:self.maxlen]
seq_ind_list.append(token_ids)
return seq_ind_list
def save_vocabulary(self, save_directory, filename_prefix=None):
if not os.path.isdir(save_directory):
os.makedirs(save_directory)
vocab_file = os.path.join(save_directory, (filename_prefix + '-' if filename_prefix else '') + 'vocab.json')
with open(vocab_file, 'w') as f:
json.dump(self.tokendict, f)
return (vocab_file,)
def save_pretrained(self, save_directory, **kwargs):
special_tokens_map_file = os.path.join(save_directory, "special_tokens_map.json")
with open(special_tokens_map_file, "w") as f:
json.dump({
"kmerlen": self.kmerlen,
"overlapping": self.overlapping,
"maxlen": self.maxlen
}, f)
vocab_files = self.save_vocabulary(save_directory)
return (special_tokens_map_file,) + vocab_files
@classmethod
def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
# Load tokenizer using the parent class method
tokenizer = super().from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs)
# Load special tokens map
special_tokens_map_file = os.path.join(pretrained_model_name_or_path, "special_tokens_map.json")
if os.path.isfile(special_tokens_map_file):
with open(special_tokens_map_file, "r") as f:
special_tokens_map = json.load(f)
tokenizer.kmerlen = special_tokens_map.get("kmerlen", 6)
tokenizer.overlapping = special_tokens_map.get("overlapping", True)
tokenizer.maxlen = special_tokens_map.get("maxlen", 400)
# Load vocabulary
vocab_file = os.path.join(pretrained_model_name_or_path, "vocab.json")
if os.path.isfile(vocab_file):
with open(vocab_file, "r") as f:
tokendict = json.load(f)
tokenizer.tokendict = tokendict
return tokenizer
|