| import json |
| import os |
|
|
| from transformers import PreTrainedTokenizer |
|
|
| _DEFAULT_VOCAB = { |
| "[UNK]": 0, |
| "[CLS]": 1, |
| "[SEP]": 2, |
| "[PAD]": 3, |
| "[MASK]": 4, |
| "A": 5, |
| "U": 6, |
| "C": 7, |
| "G": 8, |
| "N": 9, |
| } |
|
|
|
|
| class StructRFMTokenizer(PreTrainedTokenizer): |
| """Single-nucleotide tokenizer for structRFM. |
| |
| Matches the upstream preprocessing (`seq.upper().replace("T", "U")`, |
| then one token per character) and adds [CLS] ... [SEP] when |
| `add_special_tokens=True`. Characters outside the vocabulary |
| (including whitespace) map to [UNK], as in the upstream tokenizer. |
| """ |
|
|
| vocab_files_names = {"vocab_file": "vocab.json"} |
| model_input_names = ["input_ids", "attention_mask"] |
|
|
| def __init__( |
| self, |
| vocab_file=None, |
| cls_token="[CLS]", |
| sep_token="[SEP]", |
| pad_token="[PAD]", |
| mask_token="[MASK]", |
| unk_token="[UNK]", |
| **kwargs, |
| ): |
| self._vocab = dict(_DEFAULT_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__( |
| cls_token=cls_token, |
| sep_token=sep_token, |
| 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, self._vocab["[UNK]"]) |
|
|
| def _convert_id_to_token(self, index): |
| return self._ids_to_tokens.get(index, "[UNK]") |
|
|
| def convert_tokens_to_string(self, tokens): |
| return "".join(tokens) |
|
|
| 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,) |
|
|
| def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None): |
| cls = [self.cls_token_id] |
| sep = [self.sep_token_id] |
| if token_ids_1 is None: |
| return cls + token_ids_0 + sep |
| return cls + token_ids_0 + sep + token_ids_1 + sep |
|
|
| 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 |
| ) |
| mask = [1] + [0] * len(token_ids_0) + [1] |
| if token_ids_1 is not None: |
| mask += [0] * len(token_ids_1) + [1] |
| return mask |
|
|
| 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) + 2) |
| return [0] * (len(token_ids_0) + 2) + [1] * (len(token_ids_1) + 1) |
|
|
