|
|
import json |
|
|
import os |
|
|
from typing import override, Optional |
|
|
from transformers import PreTrainedTokenizer |
|
|
from Bio.Data import IUPACData |
|
|
from Bio.SeqUtils import seq3 |
|
|
from itertools import product |
|
|
|
|
|
|
|
|
class AA2CodonTokenizer(PreTrainedTokenizer): |
|
|
def __init__(self, **kwargs): |
|
|
self._aas = list(IUPACData.protein_letters_1to3.values()) |
|
|
self._codons = ["".join(p) for p in product(list("ATGC"), repeat=3)] |
|
|
|
|
|
special_tokens = { |
|
|
'bos_token': '<s>', |
|
|
'pad_token': '<pad>', |
|
|
'eos_token': '</s>', |
|
|
'unk_token': '<unk>' |
|
|
} |
|
|
|
|
|
self._vocab = self._codons + self._aas + ['*'] + list(special_tokens.values()) |
|
|
|
|
|
self._token_to_id = {token: idx for idx, token in enumerate(self._vocab)} |
|
|
self._id_to_token = {idx: token for idx, token in enumerate(self._vocab)} |
|
|
|
|
|
kwargs.update(special_tokens) |
|
|
super().__init__(**kwargs) |
|
|
|
|
|
@override |
|
|
def _tokenize(self, text: str, **kwargs) -> list[str]: |
|
|
text = text.strip() |
|
|
if text.endswith('*') and text[:-1].isupper(): |
|
|
return [str(seq3(token)) for token in list(text)[:-1]] + ['*'] |
|
|
else: |
|
|
return [text[i:i+3] for i in range(0, len(text), 3)] |
|
|
|
|
|
@override |
|
|
def _convert_token_to_id(self, token: str) -> int: |
|
|
return self._token_to_id.get(token, self.unk_token_id) |
|
|
|
|
|
@override |
|
|
def _convert_id_to_token(self, index): |
|
|
return self._id_to_token.get(index, self.unk_token) |
|
|
|
|
|
@override |
|
|
def get_vocab(self) -> dict[str, int]: |
|
|
return self._token_to_id.copy() |
|
|
|
|
|
@property |
|
|
@override |
|
|
def vocab_size(self) -> int: |
|
|
return len(self._vocab) |
|
|
|
|
|
@override |
|
|
def convert_tokens_to_string(self, tokens: list[str]) -> str: |
|
|
return "/".join(tokens) |
|
|
|
|
|
@override |
|
|
def build_inputs_with_special_tokens( |
|
|
self, token_ids_0: list[int], token_ids_1: Optional[list[int]] = None |
|
|
) -> list[int]: |
|
|
return token_ids_0 + [self.eos_token_id] |
|
|
|
|
|
@override |
|
|
def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> tuple[str]: |
|
|
filename = f"{filename_prefix}-" if filename_prefix else "" |
|
|
vocab_file = os.path.join(save_directory, f"{filename}vocab.json") |
|
|
with open(vocab_file, 'w') as f: |
|
|
json.dump(self._token_to_id, f) |
|
|
return (vocab_file, ) |
|
|
|
