Taykhoom/RNA-MSM

1import json2import os3from typing import Dict, List, Optional, Union4 5import torch6from transformers import PreTrainedTokenizer7from transformers.tokenization_utils_base import BatchEncoding8 9 10_VOCAB = {11    "<cls>":  0,12    "<pad>":  1,13    "<eos>":  2,14    "<unk>":  3,15    "A":      4,16    "G":      5,17    "C":      6,18    "U":      7,19    "X":      8,20    "N":      9,21    "-":      10,22    "<mask>": 11,23}24 25 26class RNAMSMTokenizer(PreTrainedTokenizer):27    """28    Tokenizer for RNA-MSM.29 30    Vocabulary: <cls>(0) <pad>(1) <eos>(2) <unk>(3) A(4) G(5) C(6) U(7) X(8) N(9) -(10) <mask>(11)31 32    RNA-MSM is an MSA Transformer: it always expects 3D input33    (batch, num_alignments, seqlen). This tokenizer treats each input string34    as a single-sequence MSA (1 alignment row), so the standard __call__ API:35 36        enc = tokenizer(["AGCU", "GAUC"], return_tensors="pt", padding=True)37        # enc.input_ids: (2, 1, T)  -- batch of 2 single-sequence MSAs38 39    For real MSAs (multiple aligned sequences), use encode_msa():40 41        enc = tokenizer.encode_msa([["AGCU--", "AGCUUU"]], return_tensors="pt")42        # enc["input_ids"]: (1, 2, T)  -- 1 MSA with 2 alignment rows43    """44 45    vocab_files_names = {"vocab_file": "vocab.json"}46    model_input_names = ["input_ids", "attention_mask"]47 48    def __init__(49        self,50        vocab_file=None,51        cls_token="<cls>",52        pad_token="<pad>",53        eos_token="<eos>",54        unk_token="<unk>",55        mask_token="<mask>",56        **kwargs,57    ):58        if vocab_file and os.path.isfile(vocab_file):59            with open(vocab_file) as f:60                self._vocab = json.load(f)61        else:62            self._vocab = dict(_VOCAB)63        self._ids_to_tokens = {v: k for k, v in self._vocab.items()}64        super().__init__(65            cls_token=cls_token,66            pad_token=pad_token,67            eos_token=eos_token,68            unk_token=unk_token,69            mask_token=mask_token,70            **kwargs,71        )72 73    @property74    def vocab_size(self):75        return len(self._vocab)76 77    def get_vocab(self):78        return dict(self._vocab)79 80    def _tokenize(self, text):81        return list(text)82 83    def _convert_token_to_id(self, token):84        return self._vocab.get(token, self._vocab["<unk>"])85 86    def _convert_id_to_token(self, index):87        return self._ids_to_tokens.get(index, "<unk>")88 89    def save_vocabulary(self, save_directory, filename_prefix=None):90        os.makedirs(save_directory, exist_ok=True)91        fname = (filename_prefix + "-" if filename_prefix else "") + "vocab.json"92        path = os.path.join(save_directory, fname)93        with open(path, "w") as f:94            json.dump(self._vocab, f, indent=2)95        return (path,)96 97    def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None):98        cls = [self.cls_token_id]99        if token_ids_1 is None:100            return cls + token_ids_0101        return cls + token_ids_0 + cls + token_ids_1102 103    def get_special_tokens_mask(self, token_ids_0, token_ids_1=None,104                                already_has_special_tokens=False):105        if already_has_special_tokens:106            return super().get_special_tokens_mask(107                token_ids_0, token_ids_1, already_has_special_tokens=True)108        mask = [1] + [0] * len(token_ids_0)109        if token_ids_1 is not None:110            mask += [1] + [0] * len(token_ids_1)111        return mask112 113    def create_token_type_ids_from_sequences(self, token_ids_0, token_ids_1=None):114        if token_ids_1 is None:115            return [0] * (len(token_ids_0) + 1)116        return [0] * (len(token_ids_0) + 1) + [1] * (len(token_ids_1) + 1)117 118    def __call__(119        self,120        text,121        text_pair=None,122        add_special_tokens=True,123        padding=False,124        truncation=False,125        max_length=None,126        return_tensors=None,127        **kwargs,128    ):129        """130        Tokenize one or more sequences, each treated as a 1-row MSA.131 132        text: str or List[str]133        Returns dict with input_ids of shape (batch, 1, seqlen) and134        attention_mask of shape (batch, 1, seqlen).135        """136        if isinstance(text, str):137            sequences = [text]138        else:139            sequences = list(text)140 141        encoded = []142        for seq in sequences:143            ids = self._tokenize_single(seq, add_special_tokens)144            effective_max_length = max_length145            if truncation and effective_max_length is None:146                effective_max_length = self.model_max_length147            if effective_max_length is not None and len(ids) > effective_max_length:148                if truncation:149                    ids = ids[:effective_max_length]150                else:151                    raise ValueError(152                        f"Encoded length {len(ids)} exceeds max_length "153                        f"{effective_max_length}; enable truncation."154                    )155            encoded.append(ids)156 157        if padding:158            if padding == "max_length":159                if max_length is None:160                    max_length = self.model_max_length161                max_len = max_length162            else:163                max_len = max(len(ids) for ids in encoded)164            if any(len(ids) > max_len for ids in encoded):165                raise ValueError(166                    "Cannot pad sequences that exceed the requested length."167                )168            pad_id = self.pad_token_id169            encoded = [ids + [pad_id] * (max_len - len(ids)) for ids in encoded]170 171        input_ids = [[ids] for ids in encoded]172        attention_mask = [[[1 if t != self.pad_token_id else 0 for t in ids]]173                          for ids in encoded]174 175        if return_tensors == "pt":176            input_ids = torch.tensor(input_ids, dtype=torch.long)177            attention_mask = torch.tensor(attention_mask, dtype=torch.long)178            return BatchEncoding({"input_ids": input_ids, "attention_mask": attention_mask})179 180        return BatchEncoding({"input_ids": input_ids, "attention_mask": attention_mask})181 182    def _tokenize_single(self, sequence, add_special_tokens=True):183        tokens = []184        position = 0185        special_tokens = sorted(186            self.all_special_tokens,187            key=len,188            reverse=True,189        )190        while position < len(sequence):191            special = next(192                (193                    token194                    for token in special_tokens195                    if sequence.startswith(token, position)196                ),197                None,198            )199            if special is not None:200                tokens.append(special)201                position += len(special)202            else:203                tokens.append(sequence[position])204                position += 1205        ids = [self._convert_token_to_id(t) for t in tokens]206        if add_special_tokens:207            ids = [self.cls_token_id] + ids208        return ids209 210    def encode_msa(211        self,212        msas,213        add_special_tokens=True,214        padding=False,215        return_tensors=None,216    ):217        """218        Tokenize a batch of MSAs.219 220        msas: List[List[str]]221            Each inner list is one MSA (multiple aligned sequences of equal length).222            All sequences within an MSA must have the same length.223 224        Returns dict with:225            input_ids: (batch, max_alignments, max_seqlen)226            attention_mask: (batch, max_alignments, max_seqlen)227        """228        if isinstance(msas[0], str):229            msas = [msas]230 231        for msa in msas:232            if not msa:233                raise ValueError("MSAs must contain at least one sequence.")234            if len({len(sequence) for sequence in msa}) != 1:235                raise ValueError(236                    "All sequences in an MSA must have equal aligned length."237                )238 239        max_rows = max(len(msa) for msa in msas)240        max_seqlen = max(241            len(self._tokenize_single(seq, add_special_tokens))242            for msa in msas for seq in msa243        )244        if max_seqlen > self.model_max_length:245            raise ValueError(246                f"Encoded MSA length {max_seqlen} exceeds model_max_length "247                f"{self.model_max_length}."248            )249 250        pad_id = self.pad_token_id251        batch_ids = []252        batch_mask = []253 254        for msa in msas:255            msa_ids = []256            msa_mask = []257            for seq in msa:258                ids = self._tokenize_single(seq, add_special_tokens)259                if padding:260                    pad_len = max_seqlen - len(ids)261                    mask = [1] * len(ids) + [0] * pad_len262                    ids = ids + [pad_id] * pad_len263                else:264                    mask = [1] * len(ids)265                msa_ids.append(ids)266                msa_mask.append(mask)267 268            if padding:269                pad_row = [pad_id] * max_seqlen270                pad_mask_row = [0] * max_seqlen271                while len(msa_ids) < max_rows:272                    msa_ids.append(pad_row)273                    msa_mask.append(pad_mask_row)274 275            batch_ids.append(msa_ids)276            batch_mask.append(msa_mask)277 278        if return_tensors == "pt":279            batch_ids = torch.tensor(batch_ids, dtype=torch.long)280            batch_mask = torch.tensor(batch_mask, dtype=torch.long)281            return BatchEncoding({"input_ids": batch_ids, "attention_mask": batch_mask})282 283        return BatchEncoding({"input_ids": batch_ids, "attention_mask": batch_mask})284 285    def decode(self, token_ids, skip_special_tokens=False, **kwargs):286        if isinstance(token_ids, torch.Tensor):287            token_ids = token_ids.tolist()288        tokens = [self._convert_id_to_token(i) for i in token_ids]289        if skip_special_tokens:290            special = {self.cls_token, self.pad_token, self.eos_token,291                       self.unk_token, self.mask_token}292            tokens = [t for t in tokens if t not in special]293        return "".join(tokens)294 295    def num_special_tokens_to_add(self, pair=False):296        return 1297