FishCaduceus/FishCaduceus-28L-1024

1"""Character tokenizer for Hugging Face.2 3"""4 5from typing import List, Optional, Dict, Sequence, Tuple6 7from transformers import PreTrainedTokenizer8 9 10class CaduceusTokenizer(PreTrainedTokenizer):11    model_input_names = ["input_ids"]12 13    def __init__(self,14                 model_max_length: int,15                 characters: Sequence[str] = ("A", "C", "G", "T", "N"),16                 complement_map=None,17                 bos_token="[BOS]",18                 eos_token="[SEP]",19                 sep_token="[SEP]",20                 cls_token="[CLS]",21                 pad_token="[PAD]",22                 mask_token="[MASK]",23                 unk_token="[UNK]",24                 **kwargs):25        """Character tokenizer for Hugging Face transformers.26 27        Adapted from https://huggingface.co/LongSafari/hyenadna-tiny-1k-seqlen-hf/blob/main/tokenization_hyena.py28        Args:29            model_max_length (int): Model maximum sequence length.30            characters (Sequence[str]): List of desired characters. Any character which31                is not included in this list will be replaced by a special token called32                [UNK] with id=6. Following is a list of the special tokens with33                their corresponding ids:34                    "[CLS]": 035                    "[SEP]": 136                    "[BOS]": 237                    "[MASK]": 338                    "[PAD]": 439                    "[RESERVED]": 540                    "[UNK]": 641                an id (starting at 7) will be assigned to each character.42            complement_map (Optional[Dict[str, str]]): Dictionary with string complements for each character.43        """44        if complement_map is None:45            complement_map = {"A": "T", "C": "G", "G": "C", "T": "A"}46        self.characters = characters47        self.model_max_length = model_max_length48 49        self._vocab_str_to_int = {50            "[CLS]": 0,51            "[SEP]": 1,52            "[BOS]": 2,53            "[MASK]": 3,54            "[PAD]": 4,55            "[RESERVED]": 5,56            "[UNK]": 6,57            **{ch: i + 7 for i, ch in enumerate(self.characters)},58        }59        self._vocab_int_to_str = {v: k for k, v in self._vocab_str_to_int.items()}60        add_prefix_space = kwargs.pop("add_prefix_space", False)61        padding_side = kwargs.pop("padding_side", "left")62 63        self._complement_map = {}64        for k, v in self._vocab_str_to_int.items():65            complement_id = self._vocab_str_to_int[complement_map[k]] if k in complement_map.keys() else v66            self._complement_map[self._vocab_str_to_int[k]] = complement_id67 68        super().__init__(69            bos_token=bos_token,70            eos_token=eos_token,71            sep_token=sep_token,72            cls_token=cls_token,73            pad_token=pad_token,74            mask_token=mask_token,75            unk_token=unk_token,76            add_prefix_space=add_prefix_space,77            model_max_length=model_max_length,78            padding_side=padding_side,79            **kwargs,80        )81 82    @property83    def vocab_size(self) -> int:84        return len(self._vocab_str_to_int)85 86    @property87    def complement_map(self) -> Dict[int, int]:88        return self._complement_map89 90    def _tokenize(self, text: str, **kwargs) -> List[str]:91        return list(text.upper())  # Convert all base pairs to uppercase92 93    def _convert_token_to_id(self, token: str) -> int:94        return self._vocab_str_to_int.get(token, self._vocab_str_to_int["[UNK]"])95 96    def _convert_id_to_token(self, index: int) -> str:97        return self._vocab_int_to_str[index]98 99    def convert_tokens_to_string(self, tokens):100        return "".join(tokens)  # Note: this operation has lost info about which base pairs were originally lowercase101 102    def get_special_tokens_mask(103        self,104        token_ids_0: List[int],105        token_ids_1: Optional[List[int]] = None,106        already_has_special_tokens: bool = False,107    ) -> List[int]:108        if already_has_special_tokens:109            return super().get_special_tokens_mask(110                token_ids_0=token_ids_0,111                token_ids_1=token_ids_1,112                already_has_special_tokens=True,113            )114 115        result = ([0] * len(token_ids_0)) + [1]116        if token_ids_1 is not None:117            result += ([0] * len(token_ids_1)) + [1]118        return result119 120    def build_inputs_with_special_tokens(121        self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None122    ) -> List[int]:123        sep = [self.sep_token_id]124        # cls = [self.cls_token_id]125        result = token_ids_0 + sep126        if token_ids_1 is not None:127            result += token_ids_1 + sep128        return result129 130    def get_vocab(self) -> Dict[str, int]:131        return self._vocab_str_to_int132 133    # Fixed vocabulary with no vocab file134    def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple:135        return ()136