| from transformers.tokenization_utils import AddedToken, PreTrainedTokenizer |
| from typing import List, Optional, Union |
| import os |
|
|
| MASK = "#" |
| MSA_PAD = "!" |
| UL_ALPHABET_PLUS = "ACDEFGHIKLMNPQRSTVWYBZXJOU-*#@!/[]{}" |
| MSA_AAS = "ACDEFGHIKLMNPQRSTVWYBZXJOU-" |
| GAP = "-" |
| START = "@" |
| STOP = "*" |
| SEP = "/" |
| END_AL = "]" |
| END_UL = "}" |
| START_AL = "[" |
| START_UL = "{" |
|
|
| class ProteinTokenizer(PreTrainedTokenizer): |
|
|
| def __init__( |
| self, |
| protein_alphabet: str = UL_ALPHABET_PLUS, |
| model_max_length: int = 2048, |
| pad_token=MSA_PAD, |
| mask_token=MASK, |
| all_aas=MSA_AAS, |
| gap_token=GAP, |
| bos_token=START, |
| eos_token=STOP, |
| sep_token=SEP, |
| **kwargs |
| ): |
| """Character tokenizer for Hugging Face transformers. |
| |
| model_max_length (int): Model maximum sequence length. |
| """ |
| self.alphabet = list("".join(protein_alphabet)) |
| self.all_aas = list("".join(all_aas)) |
| self.a_to_i = {u: i for i, u in enumerate(self.alphabet)} |
| self.i_to_a = {i: u for i, u in enumerate(self.alphabet)} |
| self.gap_token = gap_token |
|
|
| |
| bos_token = AddedToken(bos_token, lstrip=False, rstrip=False) if isinstance(bos_token, str) else bos_token |
| eos_token = AddedToken(eos_token, lstrip=False, rstrip=False) if isinstance(eos_token, str) else eos_token |
| sep_token = AddedToken(sep_token, lstrip=False, rstrip=False) if isinstance(sep_token, str) else sep_token |
| mask_token = AddedToken(mask_token, lstrip=False, rstrip=False) if isinstance(mask_token, str) else mask_token |
| pad_token = AddedToken(pad_token, lstrip=False, rstrip=False) if isinstance(pad_token, str) else pad_token |
| gap_token = AddedToken(gap_token, lstrip=False, rstrip=False) if isinstance(gap_token, str) else gap_token |
|
|
| super().__init__( |
| pad_token=pad_token, |
| mask_token=mask_token, |
| eos_token=eos_token, |
| bos_token=bos_token, |
| sep_token=sep_token, |
| model_max_length=model_max_length, |
| **kwargs |
| ) |
|
|
| @property |
| def vocab_size(self): |
| return len(self.alphabet) |
| |
| @property |
| def gap_token_id(self): |
| return self.convert_tokens_to_ids(self.gap_token) |
|
|
| def get_vocab(self): |
| return self.a_to_i |
|
|
| def _tokenize(self, text: str) -> List[str]: |
| return list(text) |
|
|
| def _convert_token_to_id(self, token) -> int: |
| return self.a_to_i[token] |
|
|
| def _convert_id_to_token(self, index) -> str: |
| return self.i_to_a[index] |
|
|
| def convert_tokens_to_string(self, tokens): |
| return "".join(tokens) |
|
|
| def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None): |
| result = token_ids_0 |
| if token_ids_1 is not None: |
| raise NotImplementedError("This tokenizer does not support two sequences") |
| return result |
|
|
| def get_special_tokens_mask( |
| self, |
| token_ids_0: List[int], |
| token_ids_1: Optional[List[int]] = None, |
| already_has_special_tokens: bool = False, |
| ) -> List[int]: |
| if already_has_special_tokens: |
| return super().get_special_tokens_mask( |
| token_ids_0=token_ids_0, |
| token_ids_1=token_ids_1, |
| already_has_special_tokens=True, |
| ) |
|
|
| result = [0] * len(token_ids_0) |
| if token_ids_1 is not None: |
| raise NotImplementedError("This tokenizer does not support two sequences") |
|
|
| return result |
|
|
| def create_token_type_ids_from_sequences( |
| self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None |
| ) -> List[int]: |
| """ |
| Identifies the type of token. 0 for the first sentence, 1 for the second sentence if it exists |
| """ |
|
|
| result = len(token_ids_0) * [0] |
|
|
| if token_ids_1 is not None: |
| raise NotImplementedError("This tokenizer does not support two sequences") |
| return result |
|
|
| def save_pretrained(self, save_directory: Union[str, os.PathLike], **kwargs): |
| super().save_pretrained(save_directory, **kwargs) |
| |
| def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None): |
| return () |
