diffusion-reasoning/gdsd_countdown_dream

1# coding=utf-82# Copyright 2024 The Dream team, HKUNLP Group and The HuggingFace Inc. team. All rights reserved.3#4# This code is based on Qwen's implementations in this library.5# Licensed under the Apache License, Version 2.0 (the "License");6# you may not use this file except in compliance with the License.7# You may obtain a copy of the License at8#9#     http://www.apache.org/licenses/LICENSE-2.010#11# Unless required by applicable law or agreed to in writing, software12# distributed under the License is distributed on an "AS IS" BASIS,13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.14# See the License for the specific language governing permissions and15# limitations under the License.16"""Tokenization classes for Dream."""17 18import json19import os20import unicodedata21from functools import lru_cache22from typing import Optional, Tuple23 24import regex as re25 26from transformers.tokenization_utils import AddedToken, PreTrainedTokenizer27from transformers.utils import logging28 29 30logger = logging.get_logger(__name__)31 32VOCAB_FILES_NAMES = {33    "vocab_file": "vocab.json",34    "merges_file": "merges.txt",35}36 37 38MAX_MODEL_INPUT_SIZES = {"dream/dream-tokenizer": 32768}39 40PRETOKENIZE_REGEX = r"""(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+"""41 42 43@lru_cache()44# Copied from transformers.models.gpt2.tokenization_gpt2.bytes_to_unicode45def bytes_to_unicode():46    """47    Returns list of utf-8 byte and a mapping to unicode strings. We specifically avoids mapping to whitespace/control48    characters the bpe code barfs on.49 50    The reversible bpe codes work on unicode strings. This means you need a large # of unicode characters in your vocab51    if you want to avoid UNKs. When you're at something like a 10B token dataset you end up needing around 5K for52    decent coverage. This is a significant percentage of your normal, say, 32K bpe vocab. To avoid that, we want lookup53    tables between utf-8 bytes and unicode strings.54    """55    bs = (56        list(range(ord("!"), ord("~") + 1)) + list(range(ord("¡"), ord("¬") + 1)) + list(range(ord("®"), ord("ÿ") + 1))57    )58    cs = bs[:]59    n = 060    for b in range(2**8):61        if b not in bs:62            bs.append(b)63            cs.append(2**8 + n)64            n += 165    cs = [chr(n) for n in cs]66    return dict(zip(bs, cs))67 68 69# Copied from transformers.models.gpt2.tokenization_gpt2.get_pairs70def get_pairs(word):71    """72    Return set of symbol pairs in a word.73 74    Word is represented as tuple of symbols (symbols being variable-length strings).75    """76    pairs = set()77    prev_char = word[0]78    for char in word[1:]:79        pairs.add((prev_char, char))80        prev_char = char81    return pairs82 83 84class DreamTokenizer(PreTrainedTokenizer):85    """86    Construct a Dream tokenizer. Based on byte-level Byte-Pair-Encoding.87 88    Same with GPT2Tokenizer, this tokenizer has been trained to treat spaces like parts of the tokens so a word will89    be encoded differently whether it is at the beginning of the sentence (without space) or not:90 91    ```python92    >>> from transformers import AutoTokenizer93 94    >>> tokenizer = AutoTokenizer.from_pretrained("Dream-org/Dream-v0-Base-7B", trust_remote_code=True)95    >>> tokenizer("Hello world")["input_ids"]96    [9707, 1879]97 98    >>> tokenizer(" Hello world")["input_ids"]99    [21927, 1879]100    ```101    This is expected.102 103    You should not use GPT2Tokenizer instead, because of the different pretokenization rules.104 105    This tokenizer inherits from [`PreTrainedTokenizer`] which contains most of the main methods. Users should refer to106    this superclass for more information regarding those methods.107 108    Args:109        vocab_file (`str`):110            Path to the vocabulary file.111        merges_file (`str`):112            Path to the merges file.113        errors (`str`, *optional*, defaults to `"replace"`):114            Paradigm to follow when decoding bytes to UTF-8. See115            [bytes.decode](https://docs.python.org/3/library/stdtypes.html#bytes.decode) for more information.116        unk_token (`str`, *optional*, defaults to `"<|endoftext|>"`):117            The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this118            token instead.119        bos_token (`str`, *optional*):120            The beginning of sequence token. Not applicable for this tokenizer.121        eos_token (`str`, *optional*, defaults to `"<|endoftext|>"`):122            The end of sequence token.123        pad_token (`str`, *optional*, defaults to `"<|endoftext|>"`):124            The token used for padding, for example when batching sequences of different lengths.125        clean_up_tokenization_spaces (`bool`, *optional*, defaults to `False`):126            Whether or not the model should cleanup the spaces that were added when splitting the input text during the127            tokenization process. Not applicable to this tokenizer, since tokenization does not add spaces.128        split_special_tokens (`bool`, *optional*, defaults to `False`):129            Whether or not the special tokens should be split during the tokenization process. The default behavior is130            to not split special tokens. This means that if `<|endoftext|>` is the `eos_token`, then `tokenizer.tokenize("<|endoftext|>") =131            ['<|endoftext|>`]. Otherwise, if `split_special_tokens=True`, then `tokenizer.tokenize("<|endoftext|>")` will be give `['<',132            '|', 'endo', 'ft', 'ext', '|', '>']`. This argument is only supported for `slow` tokenizers for the moment.133    """134 135    vocab_files_names = VOCAB_FILES_NAMES136    model_input_names = ["input_ids", "attention_mask"]137 138    def __init__(139        self,140        vocab_file,141        merges_file,142        errors="replace",143        unk_token="<|endoftext|>",144        bos_token=None,145        eos_token="<|endoftext|>",146        pad_token="<|endoftext|>",147        clean_up_tokenization_spaces=False,148        split_special_tokens=False,149        **kwargs,150    ):151        # Dream vocab does not contain control tokens; added tokens need to be special152        bos_token = (153            AddedToken(bos_token, lstrip=False, rstrip=False, special=True, normalized=False)154            if isinstance(bos_token, str)155            else bos_token156        )157        eos_token = (158            AddedToken(eos_token, lstrip=False, rstrip=False, special=True, normalized=False)159            if isinstance(eos_token, str)160            else eos_token161        )162        unk_token = (163            AddedToken(unk_token, lstrip=False, rstrip=False, special=True, normalized=False)164            if isinstance(unk_token, str)165            else unk_token166        )167        pad_token = (168            AddedToken(pad_token, lstrip=False, rstrip=False, special=True, normalized=False)169            if isinstance(pad_token, str)170            else pad_token171        )172 173        with open(vocab_file, encoding="utf-8") as vocab_handle:174            self.encoder = json.load(vocab_handle)175        self.decoder = {v: k for k, v in self.encoder.items()}176        self.errors = errors  # how to handle errors in decoding177        self.byte_encoder = bytes_to_unicode()178        self.byte_decoder = {v: k for k, v in self.byte_encoder.items()}179        bpe_merges = []180        with open(merges_file, encoding="utf-8") as merges_handle:181            for i, line in enumerate(merges_handle):182                line = line.strip()183                if (i == 0 and line.startswith("#version:")) or not line:184                    continue185                bpe_merges.append(tuple(line.split()))186        self.bpe_ranks = dict(zip(bpe_merges, range(len(bpe_merges))))187        # NOTE: the cache can grow without bound and will get really large for long running processes188        # (esp. for texts of language that do not use space between word, e.g. Chinese); technically189        # not a memory leak but appears as one.190        # GPT2Tokenizer has the same problem, so let's be consistent.191        self.cache = {}192 193        self.pat = re.compile(PRETOKENIZE_REGEX)194 195        if kwargs.get("add_prefix_space", False):196            logger.warning_once(197                f"{self.__class__.__name} does not support `add_prefix_space`, setting it to True has no effect."198            )199 200        super().__init__(201            errors=errors,202            bos_token=bos_token,203            eos_token=eos_token,204            pad_token=pad_token,205            unk_token=unk_token,206            clean_up_tokenization_spaces=clean_up_tokenization_spaces,207            split_special_tokens=split_special_tokens,208            **kwargs,209        )210 211    @property212    def vocab_size(self) -> int:213        return len(self.encoder)214 215    # Copied from transformers.models.gpt2.tokenization_gpt2.GPT2Tokenizer.get_vocab216    def get_vocab(self):217        return dict(self.encoder, **self.added_tokens_encoder)218 219    # Copied from transformers.models.gpt2.tokenization_gpt2.GPT2Tokenizer.bpe220    def bpe(self, token):221        if token in self.cache:222            return self.cache[token]223        word = tuple(token)224        pairs = get_pairs(word)225 226        if not pairs:227            return token228 229        while True:230            bigram = min(pairs, key=lambda pair: self.bpe_ranks.get(pair, float("inf")))231            if bigram not in self.bpe_ranks:232                break233            first, second = bigram234            new_word = []235            i = 0236            while i < len(word):237                try:238                    j = word.index(first, i)239                except ValueError:240                    new_word.extend(word[i:])241                    break242                else:243                    new_word.extend(word[i:j])244                    i = j245 246                if word[i] == first and i < len(word) - 1 and word[i + 1] == second:247                    new_word.append(first + second)248                    i += 2249                else:250                    new_word.append(word[i])251                    i += 1252            new_word = tuple(new_word)253            word = new_word254            if len(word) == 1:255                break256            else:257                pairs = get_pairs(word)258        word = " ".join(word)259        self.cache[token] = word260        return word261 262    # Copied from transformers.models.gpt2.tokenization_gpt2.GPT2Tokenizer._tokenize263    def _tokenize(self, text):264        """Tokenize a string."""265        bpe_tokens = []266        for token in re.findall(self.pat, text):267            token = "".join(268                self.byte_encoder[b] for b in token.encode("utf-8")269            )  # Maps all our bytes to unicode strings, avoiding control tokens of the BPE (spaces in our case)270            bpe_tokens.extend(bpe_token for bpe_token in self.bpe(token).split(" "))271        return bpe_tokens272 273    # Copied from transformers.models.gpt2.tokenization_gpt2.GPT2Tokenizer._convert_token_to_id274    def _convert_token_to_id(self, token):275        """Converts a token (str) in an id using the vocab."""276        return self.encoder.get(token, self.encoder.get(self.unk_token))277 278    # Copied from transformers.models.gpt2.tokenization_gpt2.GPT2Tokenizer._convert_id_to_token279    def _convert_id_to_token(self, index):280        """Converts an index (integer) in a token (str) using the vocab."""281        return self.decoder.get(index)282 283    # Copied from transformers.models.gpt2.tokenization_gpt2.GPT2Tokenizer.convert_tokens_to_string284    def convert_tokens_to_string(self, tokens):285        """Converts a sequence of tokens (string) in a single string."""286        text = "".join(tokens)287        text = bytearray([self.byte_decoder[c] for c in text]).decode("utf-8", errors=self.errors)288        return text289 290    def decode(291        self,292        token_ids,293        skip_special_tokens: bool = False,294        clean_up_tokenization_spaces: Optional[bool] = False,295        spaces_between_special_tokens: bool = False,296        **kwargs,297    ) -> str:298        # `spaces_between_special_tokens` defaults to True for _decode in slow tokenizers299        # and cannot be configured elsewhere, but it should default to False for DreamTokenizer300        return super().decode(301            token_ids,302            skip_special_tokens=skip_special_tokens,303            clean_up_tokenization_spaces=clean_up_tokenization_spaces,304            spaces_between_special_tokens=spaces_between_special_tokens,305            **kwargs,306        )307 308    # Copied from transformers.models.gpt2.tokenization_gpt2.GPT2Tokenizer.save_vocabulary309    def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:310        if not os.path.isdir(save_directory):311            logger.error(f"Vocabulary path ({save_directory}) should be a directory")312            return313        vocab_file = os.path.join(314            save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]315        )316        merge_file = os.path.join(317            save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["merges_file"]318        )319 320        with open(vocab_file, "w", encoding="utf-8") as f:321            f.write(json.dumps(self.encoder, indent=2, sort_keys=True, ensure_ascii=False) + "\n")322 323        index = 0324        with open(merge_file, "w", encoding="utf-8") as writer:325            writer.write("#version: 0.2\n")326            for bpe_tokens, token_index in sorted(self.bpe_ranks.items(), key=lambda kv: kv[1]):327                if index != token_index:328                    logger.warning(329                        f"Saving vocabulary to {merge_file}: BPE merge indices are not consecutive."330                        " Please check that the tokenizer is not corrupted!"331                    )332                    index = token_index333                writer.write(" ".join(bpe_tokens) + "\n")334                index += 1335 336        return vocab_file, merge_file337 338    def prepare_for_tokenization(self, text, **kwargs):339        text = unicodedata.normalize("NFC", text)340        return (text, kwargs)