File size: 1,779 Bytes
5ae0a0b | 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 | import os
import re
from transformers import PreTrainedTokenizer
class GenoJEPATokenizer(PreTrainedTokenizer):
def __init__(self, **kwargs):
self.pad_token = "N"
self.base_chars = ["A", "T", "C", "G"]
self.special_tokens = [self.pad_token]
self.vocab = {token: i for i, token in enumerate(self.base_chars + self.special_tokens)}
self.ids_to_tokens = {i: token for token, i in self.vocab.items()}
self.dna_pattern = re.compile(r"[ATCG]")
self.pad_token_id = self._convert_token_to_id(self.pad_token)
super().__init__(**kwargs)
@property
def vocab_size(self):
return len(self.vocab)
def get_vocab(self):
return dict(self.vocab)
def _convert_token_to_id(self, token):
return self.vocab.get(token, self.vocab[self.pad_token])
def _convert_id_to_token(self, index):
return self.ids_to_tokens.get(index, self.pad_token)
def _tokenize(self, seq, **kwargs):
tokens = []
pos = 0
while pos < len(seq):
dna_match = self.dna_pattern.match(seq, pos)
if dna_match:
dna_seq = dna_match.group()
tokens.append(dna_seq)
pos = dna_match.end()
else:
tokens.append(self.pad_token)
pos += 1
return tokens
def save_vocabulary(self, save_directory, filename_prefix=None):
vocab_file = os.path.join(save_directory, (filename_prefix + "-" if filename_prefix else "") + "vocab.txt")
with open(vocab_file, "w", encoding="utf-8") as writer:
for token, token_index in sorted(self.vocab.items(), key=lambda kv: kv[1]):
writer.write(token + "\n")
return (vocab_file,)
|