codingutf8 2022 meta and the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache orglicenseslicense2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license tokenization classes for esm import os from typing import list optional union from tokenizationutils import pretrainedtokenizer from tokenizationutilsbase import addedtoken from utils import logging logger logging getloggername vocabfilesnames vocabfile vocab txt pretrainedvocabfilesmap vocabfile facebookesm2t68mur50d https huggingface cofacebookesm2t68mur50dresolvemainvocab txt facebookesm2t1235mur50d https huggingface cofacebookesm2t1235mur50dresolvemainvocab txt pretrainedpositionalembeddingssizes facebookesm2t68mur50d 1024 facebookesm2t1235mur50d 1024 def loadvocabfilevocabfile with openvocabfile r as f lines f read splitlines return l strip for l in lines class esmtokenizerpretrainedtokenizer vocabfilesnames vocabfilesnames pretrainedvocabfilesmap pretrainedvocabfilesmap maxmodelinputsizes pretrainedpositionalembeddingssizes modelinputnames inputids attentionmask def init self vocabfile unktokenunk clstokencls padtokenpad masktokenmask eostokeneos kwargs self alltokens loadvocabfilevocabfile self idtotoken dictenumerateself alltokens self tokentoid tok ind for ind tok in enumerateself alltokens super init unktokenunktoken clstokenclstoken padtokenpadtoken masktokenmasktoken eostokeneostoken kwargs todo all the tokens are added but they are also part of the vocab bit strange none of them are special but they all need special splitting self uniquenosplittokens self alltokens self updatetrieself uniquenosplittokens def convertidtotokenself index int str return self idtotoken getindex self unktoken def converttokentoidself token str int return self tokentoid gettoken self tokentoid getself unktoken def tokenizeself text kwargs return text split def getvocabsizeself withaddedtokensfalse return lenself idtotoken def getvocabself return token i for i token in enumerateself alltokens def tokentoidself token str int return self tokentoid gettoken self tokentoid getself unktoken def idtotokenself index int str return self idtotoken getindex self unktoken def buildinputswithspecialtokens self tokenids0 listint tokenids1 optionallistint none listint cls self clstokenid sep self eostokenid no sep token in esm vocabulary if tokenids1 is none if self eostokenid is none return cls tokenids0 else return cls tokenids0 sep elif self eostokenid is none raise valueerrorcannot tokenize multiple sequences when eos token is not set return cls tokenids0 sep tokenids1 sep multiple inputs always have an eos token def getspecialtokensmask self tokenids0 list tokenids1 optionallist none alreadyhasspecialtokens bool false listint if alreadyhasspecialtokens if tokenids1 is not none raise valueerror you should not supply a second sequence if the provided sequence of ids is already formatted with special tokens for the model return 1 if token in self allspecialids else 0 for token in tokenids0 mask 1 0 lentokenids0 1 if tokenids1 is not none mask 0 lentokenids1 1 return mask def savevocabularyself savedirectory filenameprefix vocabfile os path joinsavedirectory filenameprefix if filenameprefix else vocab txt with openvocabfile w as f f writen joinself alltokens return vocabfile property def vocabsizeself int return self getvocabsizewithaddedtokensfalse def addtokensself newtokens unionliststr listaddedtoken specialtokens bool false int return super addtokensnewtokens specialtokenstrue coding utf 8 2022 meta and the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache org licenses license 2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license tokenization classes for esm constructs an esm tokenizer todo all the tokens are added but they are also part of the vocab bit strange none of them are special but they all need special splitting no sep token in esm vocabulary multiple inputs always have an eos token retrieves sequence ids from a token list that has no special tokens added this method is called when adding special tokens using the tokenizer prepare_for_model or encode_plus methods args token_ids_0 list int list of ids of the first sequence token_ids_1 list int optional list of ids of the second sequence already_has_special_tokens bool optional defaults to false whether or not the token list is already formatted with special tokens for the model returns a list of integers in the range 0 1 1 for a special token 0 for a sequence token
import os
from typing import List, Optional, Union

from ...tokenization_utils import PreTrainedTokenizer
from ...tokenization_utils_base import AddedToken
from ...utils import logging


logger = logging.get_logger(__name__)

VOCAB_FILES_NAMES = {"vocab_file": "vocab.txt"}

PRETRAINED_VOCAB_FILES_MAP = {
    "vocab_file": {
        "facebook/esm2_t6_8M_UR50D": "https://huggingface.co/facebook/esm2_t6_8M_UR50D/resolve/main/vocab.txt",
        "facebook/esm2_t12_35M_UR50D": "https://huggingface.co/facebook/esm2_t12_35M_UR50D/resolve/main/vocab.txt",
    },
}

PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES = {
    "facebook/esm2_t6_8M_UR50D": 1024,
    "facebook/esm2_t12_35M_UR50D": 1024,
}


def load_vocab_file(vocab_file):
    with open(vocab_file, "r") as f:
        lines = f.read().splitlines()
        return [l.strip() for l in lines]


class EsmTokenizer(PreTrainedTokenizer):
    

    vocab_files_names = VOCAB_FILES_NAMES
    pretrained_vocab_files_map = PRETRAINED_VOCAB_FILES_MAP
    max_model_input_sizes = PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
    model_input_names = ["input_ids", "attention_mask"]

    def __init__(
        self,
        vocab_file,
        unk_token="<unk>",
        cls_token="<cls>",
        pad_token="<pad>",
        mask_token="<mask>",
        eos_token="<eos>",
        **kwargs,
    ):
        self.all_tokens = load_vocab_file(vocab_file)
        self._id_to_token = dict(enumerate(self.all_tokens))
        self._token_to_id = {tok: ind for ind, tok in enumerate(self.all_tokens)}
        super().__init__(
            unk_token=unk_token,
            cls_token=cls_token,
            pad_token=pad_token,
            mask_token=mask_token,
            eos_token=eos_token,
            **kwargs,
        )

        
        

        self.unique_no_split_tokens = self.all_tokens
        self._update_trie(self.unique_no_split_tokens)

    def _convert_id_to_token(self, index: int) -> str:
        return self._id_to_token.get(index, self.unk_token)

    def _convert_token_to_id(self, token: str) -> int:
        return self._token_to_id.get(token, self._token_to_id.get(self.unk_token))

    def _tokenize(self, text, **kwargs):
        return text.split()

    def get_vocab_size(self, with_added_tokens=False):
        return len(self._id_to_token)

    def get_vocab(self):
        return {token: i for i, token in enumerate(self.all_tokens)}

    def token_to_id(self, token: str) -> int:
        return self._token_to_id.get(token, self._token_to_id.get(self.unk_token))

    def id_to_token(self, index: int) -> str:
        return self._id_to_token.get(index, self.unk_token)

    def build_inputs_with_special_tokens(
        self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
    ) -> List[int]:
        cls = [self.cls_token_id]
        sep = [self.eos_token_id]  
        if token_ids_1 is None:
            if self.eos_token_id is None:
                return cls + token_ids_0
            else:
                return cls + token_ids_0 + sep
        elif self.eos_token_id is None:
            raise ValueError("Cannot tokenize multiple sequences when EOS token is not set!")
        return cls + token_ids_0 + sep + token_ids_1 + sep  

    def get_special_tokens_mask(
        self, token_ids_0: List, token_ids_1: Optional[List] = None, already_has_special_tokens: bool = False
    ) -> List[int]:
        
        if already_has_special_tokens:
            if token_ids_1 is not None:
                raise ValueError(
                    "You should not supply a second sequence if the provided sequence of "
                    "ids is already formatted with special tokens for the model."
                )

            return [1 if token in self.all_special_ids else 0 for token in token_ids_0]
        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 save_vocabulary(self, save_directory, filename_prefix):
        vocab_file = os.path.join(save_directory, (filename_prefix + "-" if filename_prefix else "") + "vocab.txt")
        with open(vocab_file, "w") as f:
            f.write("\n".join(self.all_tokens))
        return (vocab_file,)

    @property
    def vocab_size(self) -> int:
        return self.get_vocab_size(with_added_tokens=False)

    def _add_tokens(self, new_tokens: Union[List[str], List[AddedToken]], special_tokens: bool = False) -> int:
        return super()._add_tokens(new_tokens, special_tokens=True)