| import json |
| import os |
| from transformers import PreTrainedTokenizer |
|
|
| VOCAB = {"<pad>": 0, "<mask>": 1, "A": 2, "U": 3, "G": 4, "C": 5} |
|
|
|
|
| class RNABertTokenizer(PreTrainedTokenizer): |
| vocab_files_names = {"vocab_file": "vocab.json"} |
| model_input_names = ["input_ids", "attention_mask"] |
|
|
| def __init__( |
| self, |
| vocab_file=None, |
| pad_token="<pad>", |
| mask_token="<mask>", |
| unk_token="<pad>", |
| **kwargs, |
| ): |
| self._vocab = dict(VOCAB) |
| if vocab_file and os.path.isfile(vocab_file): |
| with open(vocab_file) as f: |
| self._vocab = json.load(f) |
| self._ids_to_tokens = {v: k for k, v in self._vocab.items()} |
| super().__init__( |
| pad_token=pad_token, |
| mask_token=mask_token, |
| unk_token=unk_token, |
| **kwargs, |
| ) |
|
|
| @property |
| def vocab_size(self): |
| return len(self._vocab) |
|
|
| def get_vocab(self): |
| return dict(self._vocab) |
|
|
| def _tokenize(self, text): |
| return list(text.upper().replace("T", "U")) |
|
|
| def _convert_token_to_id(self, token): |
| return self._vocab.get(token, 0) |
|
|
| def _convert_id_to_token(self, index): |
| return self._ids_to_tokens.get(index, "<pad>") |
|
|
| def save_vocabulary(self, save_directory, filename_prefix=None): |
| os.makedirs(save_directory, exist_ok=True) |
| fname = (filename_prefix + "-" if filename_prefix else "") + "vocab.json" |
| path = os.path.join(save_directory, fname) |
| with open(path, "w") as f: |
| json.dump(self._vocab, f, indent=2) |
| return (path,) |
|
|
| @property |
| def cls_token_id(self): |
| return self.pad_token_id |
|
|
| @property |
| def eos_token_id(self): |
| return self.pad_token_id |
|
|
| def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None): |
| if token_ids_1 is None: |
| return token_ids_0 |
| return token_ids_0 + token_ids_1 |
|
|
| 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, True) |
| return [0] * len(token_ids_0) + ([0] * len(token_ids_1) if token_ids_1 else []) |
|
|
| def create_token_type_ids_from_sequences(self, token_ids_0, token_ids_1=None): |
| if token_ids_1 is None: |
| return [0] * len(token_ids_0) |
| return [0] * len(token_ids_0) + [0] * len(token_ids_1) |
|
|
