| import itertools |
| import os |
| import json |
| import re |
| from typing import List, Optional, Tuple |
| from transformers import PreTrainedTokenizer |
|
|
| class DNAKmerTokenizer(PreTrainedTokenizer): |
| def __init__(self, k, **kwargs): |
| self.k = k |
| self.special_tokens = [ |
| "<oov>", |
| "<s>", |
| "</s>", |
| "<pad>", |
| "<mask>", |
| "<bog>", |
| "<eog>", |
| "<bok>", |
| "<eok>", |
| "<+>", |
| "<->", |
| "<cds>", |
| "<pseudo>", |
| "<tRNA>", |
| "<rRNA>", |
| "<ncRNA>", |
| "<miscRNA>", |
| "<mam>", |
| "<vrt>", |
| "<inv>", |
| "<pln>", |
| "<fng>", |
| "<prt>", |
| "<arc>", |
| "<bct>", |
| "<mit>", |
| "<plt>", |
| "<plm>", |
| "<vir>", |
| "<sp0>", |
| "<sp1>", |
| "<sp2>", |
| "<crispr_spacer>", |
| "<crispr_repeat>", |
| "<cas1>", |
| "<cas2>", |
| "<tracrrna>", |
| "<cas5>", |
| "<cas3>", |
| "<cas4>", |
| "<cas9>", |
| "<cas7>", |
| "<cas8c>", |
| "<cas6>", |
| "<csm3gr7>", |
| "<csn2>", |
| "<cas10>", |
| "<cas7b>", |
| "<cas6e>", |
| "<cas8e>", |
| "<cas12f>", |
| "<cse2gr11>", |
| "<csx1>", |
| "<csm2gr11>", |
| "<csm4>", |
| "<wyl>", |
| "<cas12a>", |
| "<csm6>", |
| "<deddh>", |
| "<csm5>", |
| "<casr>", |
| "<cas8b1>", |
| "<csx19>", |
| "<csx20>", |
| "<csm5gr7>", |
| "<cas6f>", |
| "<cas8b2>", |
| "<cas5f>", |
| "<rt>", |
| "<cas7f>", |
| "<cas3-cas2>", |
| "<primpol>", |
| "<cas8f>", |
| "<cysh>", |
| "<cas3hd>", |
| "<tnib>", |
| "<csx10gr5>", |
| "<cas8a1>", |
| "<csa3>", |
| "<recd>", |
| "<cmr1gr7>", |
| "<cmr4>", |
| "<cmr6gr7>", |
| "<cmr3gr5>", |
| "<cmr5gr11>", |
| "<cas8b6>", |
| "<csb2>", |
| "<cora>", |
| "<csm4gr5>", |
| "<abieii>", |
| "<can2>", |
| "<cas13d>", |
| "<csb1gr7>", |
| "<iscb-hnh>", |
| "<pd>", |
| "<tnpa>", |
| "<cse2>", |
| "<csb3>", |
| "<csm3>", |
| "<cas13b>", |
| "<unk>", |
| "<csx16>", |
| "<tpr>", |
| "<dhh>", |
| "<2og>", |
| "<cas12m>", |
| "<mem>", |
| "<csf4>", |
| "<hearo>", |
| "<tn7>", |
| "<tniq>", |
| "<csf2>", |
| "<csf3>", |
| "<csf1>", |
| "<cas8b4>", |
| "<tnsd>", |
| "<heat>", |
| "<csx17>", |
| "<cas8u1>", |
| "<csx3>", |
| "<htpx>", |
| "<cas12b>", |
| "<csm2>", |
| "<cas10d>", |
| "<csc2>", |
| "<cmr3>", |
| "<cmr5>", |
| "<csc1gr5>", |
| "<gramp>", |
| "<cmr6>", |
| "<cas8b12>", |
| "<cas11b>", |
| "<cas12c>", |
| "<cas8a4>", |
| "<tnsb>", |
| "<nyn>", |
| "<iscb-nterm>", |
| "<cas8b3>", |
| "<cas8a2>", |
| "<cas5u>", |
| "<csx27>", |
| "<csx21>", |
| "<csx23>", |
| "<tm>", |
| "<cas3d>", |
| "<cas12lambda>", |
| "<tnsc>", |
| "<cas8b5>", |
| "<stand>", |
| "<st>", |
| "<iscb-ruvciii-cterm>", |
| "<cas11>", |
| "<cas11d2>", |
| "<cas12j>", |
| "<cas12d>", |
| "<cas8b8>", |
| "<cmr1>", |
| "<cas12k>", |
| "<cas12g>", |
| "<cas13f>", |
| "<cas8b10>", |
| "<cas13i>", |
| "<toprim>", |
| "<cas12e>", |
| ] |
| self.kmers = [ |
| "".join(kmer) for kmer in itertools.product("ATCG", repeat=self.k) |
| ] |
| self.vocab = { |
| token: i for i, token in enumerate(self.special_tokens + self.kmers) |
| } |
| self.ids_to_tokens = {v: k for k, v in self.vocab.items()} |
| self.special_token_pattern = re.compile( |
| "|".join(re.escape(token) for token in self.special_tokens) |
| ) |
| self.dna_pattern = re.compile(f"[A-Z]{{{self.k}}}|[A-Z]+") |
| self.bos_token = "<s>" |
| self.eos_token = "</s>" |
| self.bos_token_id = self._convert_token_to_id(self.bos_token) |
| self.eos_token_id = self._convert_token_to_id(self.eos_token) |
| super().__init__(**kwargs) |
|
|
| @property |
| def vocab_size(self): |
| return len(self.vocab) |
|
|
| def get_vocab(self): |
| return dict(self.vocab) |
|
|
| def _tokenize(self, text, **kwargs) -> List[str]: |
| tokens = [] |
| pos = 0 |
| while pos < len(text): |
| special_match = self.special_token_pattern.match(text, pos) |
| if special_match: |
| tokens.append(special_match.group()) |
| pos = special_match.end() |
| else: |
| dna_match = self.dna_pattern.match(text, pos) |
| if dna_match: |
| dna_seq = dna_match.group() |
| tokens.append(dna_seq) |
| pos = dna_match.end() |
| else: |
| tokens.append(text[pos]) |
| pos += 1 |
| return tokens |
|
|
| def _convert_token_to_id(self, token: str) -> int: |
| return self.vocab.get(token, self.vocab["<oov>"]) |
|
|
| def _convert_id_to_token(self, index: int) -> str: |
| return self.ids_to_tokens.get(index, "<oov>") |
|
|
| def convert_tokens_to_string(self, tokens: List[str]) -> str: |
| 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.bos_token_id] + token_ids_0 + [self.eos_token_id] |
| return [self.bos_token_id] + token_ids_0 + [self.eos_token_id] + token_ids_1 + [self.eos_token_id] |
|
|
| def get_special_tokens_mask( |
| self, token_ids_0, token_ids_1=None, already_has_special_tokens=False |
| ): |
| if already_has_special_tokens: |
| return super().get_special_tokens_mask( |
| token_ids_0, token_ids_1, already_has_special_tokens=True |
| ) |
| if token_ids_1 is None: |
| return [1] + ([0] * len(token_ids_0)) + [1] |
| return [1] + ([0] * len(token_ids_0)) + [1] + ([0] * len(token_ids_1)) + [1] |
|
|
| def prepare_for_model(self, *args, **kwargs): |
| encoding = super().prepare_for_model(*args, **kwargs) |
| if "token_type_ids" in encoding: |
| del encoding["token_type_ids"] |
| return encoding |
|
|
| def save_vocabulary( |
| self, save_directory: str, filename_prefix: Optional[str] = None |
| ) -> Tuple[str]: |
| import os |
|
|
| 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,) |
| |
| def save_pretrained(self, save_directory: str, **kwargs): |
| vocab_files = super().save_pretrained(save_directory, **kwargs) |
| tokenizer_config_path = os.path.join(save_directory, "tokenizer_config.json") |
| |
| |
| if os.path.exists(tokenizer_config_path): |
| with open(tokenizer_config_path, "r", encoding="utf-8") as f: |
| config = json.load(f) |
| else: |
| config = {} |
| |
| |
| config.update({ |
| "auto_map": { |
| "AutoTokenizer": [ |
| "tokenizer.DNAKmerTokenizer", |
| None |
| ] |
| }, |
| }) |
| |
| |
| config.update({ |
| "k": self.k |
| }) |
| |
| |
| with open(tokenizer_config_path, "w", encoding="utf-8") as f: |
| json.dump(config, f, ensure_ascii=False, indent=2) |
| |
| return vocab_files |
|
|
