Upload 4 files

special_tokens_map.json

@@ -0,0 +1,5 @@
+{
+  "bos_token": "<bos>",
+  "eos_token": "<bos>",
+  "pad_token": "<pad>"
+}

tokenization_bolmo.py

@@ -0,0 +1,378 @@
+from dataclasses import dataclass, field
+from functools import lru_cache
+from typing import Optional, Union
+from transformers import AutoTokenizer
+from transformers.tokenization_utils import PreTrainedTokenizer
+
+# Source: https://github.com/openai/gpt-2/blob/master/src/encoder.py#L9
+# Also implemented in https://docs.rs/tokenizers/latest/src/tokenizers/pre_tokenizers/byte_level.rs.html#13-39
+_CHARS_TO_BYTES = {
+    "Ā": 0, "ā": 1, "Ă": 2, "ă": 3, "Ą": 4, "ą": 5, "Ć": 6, "ć": 7, "Ĉ": 8,
+    "ĉ": 9, "Ċ": 10, "ċ": 11, "Č": 12, "č": 13, "Ď": 14, "ď": 15, "Đ": 16,
+    "đ": 17, "Ē": 18, "ē": 19, "Ĕ": 20, "ĕ": 21, "Ė": 22, "ė": 23, "Ę": 24,
+    "ę": 25, "Ě": 26, "ě": 27, "Ĝ": 28, "ĝ": 29, "Ğ": 30, "ğ": 31, "Ġ": 32,
+    "!": 33, '"': 34, "#": 35, "$": 36, "%": 37, "&": 38, "'": 39, "(": 40,
+    ")": 41, "*": 42, "+": 43, ",": 44, "-": 45, ".": 46, "/": 47, "0": 48,
+    "1": 49, "2": 50, "3": 51, "4": 52, "5": 53, "6": 54, "7": 55, "8": 56,
+    "9": 57, ":": 58, ";": 59, "<": 60, "=": 61, ">": 62, "?": 63, "@": 64, 
+    "A": 65, "B": 66, "C": 67, "D": 68, "E": 69, "F": 70, "G": 71, "H": 72,
+    "I": 73, "J": 74, "K": 75, "L": 76, "M": 77, "N": 78, "O": 79, "P": 80,
+    "Q": 81, "R": 82, "S": 83, "T": 84, "U": 85, "V": 86, "W": 87, "X": 88,
+    "Y": 89, "Z": 90, "[": 91, "\\": 92, "]": 93, "^": 94, "_": 95, "`": 96, 
+    "a": 97, "b": 98, "c": 99, "d": 100, "e": 101, "f": 102, "g": 103,
+    "h": 104, "i": 105, "j": 106, "k": 107, "l": 108, "m": 109, "n": 110,
+    "o": 111, "p": 112, "q": 113, "r": 114, "s": 115, "t": 116, "u": 117,
+    "v": 118, "w": 119, "x": 120, "y": 121, "z": 122, "{": 123, "|": 124,
+    "}": 125, "~": 126, "ġ": 127, "Ģ": 128, "ģ": 129, "Ĥ": 130, "ĥ": 131,
+    "Ħ": 132, "ħ": 133, "Ĩ": 134, "ĩ": 135, "Ī": 136, "ī": 137, "Ĭ": 138,
+    "ĭ": 139, "Į": 140, "į": 141, "İ": 142, "ı": 143, "IJ": 144, "ij": 145,
+    "Ĵ": 146, "ĵ": 147, "Ķ": 148, "ķ": 149, "ĸ": 150, "Ĺ": 151, "ĺ": 152,
+    "Ļ": 153, "ļ": 154, "Ľ": 155, "ľ": 156, "Ŀ": 157, "ŀ": 158, "Ł": 159,
+    "ł": 160, "¡": 161, "¢": 162, "£": 163, "¤": 164, "¥": 165, "¦": 166,
+    "§": 167, "¨": 168, "©": 169, "ª": 170, "«": 171, "¬": 172, "Ń": 173,
+    "®": 174, "¯": 175, "°": 176, "±": 177, "²": 178, "³": 179, "´": 180,
+    "µ": 181, "¶": 182, "·": 183, "¸": 184, "¹": 185, "º": 186, "»": 187,
+    "¼": 188, "½": 189, "¾": 190, "¿": 191, "À": 192, "Á": 193, "Â": 194,
+    "Ã": 195, "Ä": 196, "Å": 197, "Æ": 198, "Ç": 199, "È": 200, "É": 201,
+    "Ê": 202, "Ë": 203, "Ì": 204, "Í": 205, "Î": 206, "Ï": 207, "Ð": 208,
+    "Ñ": 209, "Ò": 210, "Ó": 211, "Ô": 212, "Õ": 213, "Ö": 214, "×": 215,
+    "Ø": 216, "Ù": 217, "Ú": 218, "Û": 219, "Ü": 220, "Ý": 221, "Þ": 222,
+    "ß": 223, "à": 224, "á": 225, "â": 226, "ã": 227, "ä": 228, "å": 229,
+    "æ": 230, "ç": 231, "è": 232, "é": 233, "ê": 234, "ë": 235, "ì": 236,
+    "í": 237, "î": 238, "ï": 239, "ð": 240, "ñ": 241, "ò": 242, "ó": 243,
+    "ô": 244, "õ": 245, "ö": 246, "÷": 247, "ø": 248, "ù": 249, "ú": 250,
+    "û": 251, "ü": 252, "ý": 253, "þ": 254, "ÿ": 255,
+}
+_BYTES_TO_CHARS = {v: k for k, v in _CHARS_TO_BYTES.items()}
+
+def _bytes_to_chars(byte_sequence: bytes) -> str:
+    return "".join(_BYTES_TO_CHARS[byte] for byte in byte_sequence)
+
+def _chars_to_bytes(char_sequence: str) -> list:
+    return list(bytes(_CHARS_TO_BYTES[char] for char in char_sequence))
+
+@dataclass
+class BolmoTokenizerConfig:
+    vocab_size: int
+    bos_token_id: int
+    pad_token_id: int
+    eos_token_id: int
+    bpe_token_end_id: int
+    special_tokens: list[str] = field(default_factory=lambda: [])
+    special_tokens_first: bool = True
+    original_identifier: Optional[str] = None
+
+
+    @classmethod
+    def bolmo(cls) -> "BolmoTokenizerConfig":
+        special_tokens = [
+            "<pad>",
+            "<bos>",
+            "<eos>",
+            "<bpe_token_end>",
+        ]
+
+        return cls(
+            # *2 to accomodate fused boundary tokens
+            vocab_size=(len(special_tokens) + 256) * 2,
+            special_tokens=special_tokens,
+            bos_token_id=special_tokens.index("<bos>"),
+            pad_token_id=special_tokens.index("<pad>"),
+            eos_token_id=special_tokens.index("<bos>"),
+            bpe_token_end_id=special_tokens.index("<bpe_token_end>"),
+            original_identifier="allenai/dolma2-tokenizer",
+        )
+    
+    def build(self):
+        return BolmoTokenizer(tokenizer_config=self)
+
+
+class BolmoTokenizer(PreTrainedTokenizer):
+    TOKEN_ID_KEY = -1
+
+    def __init__(self, **kwargs):
+        tokenizer_config = kwargs.pop("tokenizer_config", BolmoTokenizerConfig.bolmo())
+
+        self.config = tokenizer_config
+        self.hf_tokenizer = AutoTokenizer.from_pretrained(tokenizer_config.original_identifier)
+        if self.config.special_tokens_first:
+            self.offset = len(tokenizer_config.special_tokens)
+            self.special_tokens_offset = 0
+        else:
+            self.offset = 0
+            self.special_tokens_offset = self.config.vocab_size - len(tokenizer_config.special_tokens)
+
+        self.byte_sequences = {}
+    
+        for key, value in self.hf_tokenizer.get_vocab().items():
+            if key in self.config.special_tokens:
+                byte_sequence = [self.special_tokens_offset + self.config.special_tokens.index(key)]
+            elif value == self.hf_tokenizer.eos_token_id and self.eos_token_id is not None:
+                byte_sequence = [self.eos_token_id]
+            elif value == self.hf_tokenizer.bos_token_id and self.bos_token_id is not None:
+                byte_sequence = [self.bos_token_id]
+            elif value == self.hf_tokenizer.pad_token_id and self.pad_token_id is not None:
+                byte_sequence = [self.pad_token_id]
+            else:
+                byte_sequence = [self.offset + i for i in _chars_to_bytes(key)]
+
+            assert self.byte_sequences.get(value) is None
+            self.byte_sequences[value] = byte_sequence
+
+        self.byte_trie = {}
+
+        for token_id, byte_sequence in self.byte_sequences.items():
+            current_dict = self.byte_trie
+            for byte in byte_sequence[::-1]: # retrieved from the back so store in reverse order
+                if byte not in current_dict:
+                    current_dict[byte] = {}
+                current_dict = current_dict[byte]
+            current_dict[BolmoTokenizer.TOKEN_ID_KEY] = token_id
+
+        self.add_bos_token = True
+        self.add_eos_token = False
+        self.padding_side = "left" # for generate
+
+        super().__init__(
+            bos_token=self.config.special_tokens[self.config.bos_token_id],
+            eos_token=self.config.special_tokens[self.config.eos_token_id],
+            pad_token=self.config.special_tokens[self.config.pad_token_id],
+            extra_ids=0,
+        )
+
+    @property
+    def bos_token_id(self):
+        return self.config.bos_token_id
+
+    @property
+    def eos_token_id(self):
+        return self.config.eos_token_id
+    
+    @property
+    def pad_token_id(self):
+        return self.config.pad_token_id
+
+    @property
+    def bpe_token_end_id(self):
+        return self.config.bpe_token_end_id 
+
+    @property
+    def vocab_size(self):
+        return self.config.vocab_size
+
+    def _convert_id_to_token(self, index):
+        if index < self.offset:
+            return self.config.special_tokens[index - self.special_tokens_offset]
+
+        if index >= self.offset + 256 and index < self.offset * 2 + 256:
+            # special token with fused boundary
+            return self.config.special_tokens[index - self.offset - 256] + "b"
+
+        return _BYTES_TO_CHARS[index - self.offset - 256 - self.offset] + "b" if index >= self.offset + 256 else _BYTES_TO_CHARS[index - self.offset]
+
+    def _convert_token_to_id(self, token):
+        if token in self.config.special_tokens:
+            return self.config.special_tokens.index(token)
+
+        if token in [x + "b" for x in self.config.special_tokens]:
+            # special token with fused boundary
+            return 256 + self.config.special_tokens.index(token[:-1])
+
+        if len(token) > 1 and token[-1] == "b":
+            return self.offset + 256 + _CHARS_TO_BYTES[token[0]]
+        else:
+            return self.offset + _CHARS_TO_BYTES[token]
+
+    def get_vocab(self):
+        vocab = {self.convert_ids_to_tokens(i): i for i in range(self.vocab_size)}
+        return vocab
+
+    def expand_byte_ids(self, byte_ids: list[int], n_last: Optional[int] = None) -> list[int]:
+        # search in the byte tree for the longest matching token at every byte position
+        expanded_ids = []
+        for i in range(len(byte_ids)):
+            if n_last is not None and i < len(byte_ids) - n_last:
+                continue
+
+            current_dict = self.byte_trie
+            current_expansion = None
+
+            for i in range(i, -1, -1):
+                byte = byte_ids[i]
+
+                if byte == self.bpe_token_end_id:
+                    # skip bpe token end markers, needed for generation
+                    continue
+
+                if byte >= self.offset + 256:
+                    # ignore fused boundary
+                    byte -= self.offset + 256
+
+                try:
+                    current_dict = current_dict[byte]
+                    if BolmoTokenizer.TOKEN_ID_KEY in current_dict:
+                        current_expansion = current_dict[BolmoTokenizer.TOKEN_ID_KEY]
+                except KeyError:
+                    assert current_expansion is not None
+                    break
+
+            expanded_ids.append(current_expansion)
+
+        return expanded_ids
+
+    # Copied from transformers.models.llama.tokenization_llama.LlamaTokenizer.build_inputs_with_special_tokens
+    def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None):
+        bos_token_id = [self.bos_token_id] if self.add_bos_token else []
+        eos_token_id = [self.eos_token_id] if self.add_eos_token else []
+
+        output = bos_token_id + token_ids_0 + eos_token_id
+
+        if token_ids_1 is not None:
+            output = output + bos_token_id + token_ids_1 + eos_token_id
+
+        return output
+
+    # Copied from transformers.models.llama.tokenization_llama.LlamaTokenizer.get_special_tokens_mask
+    def get_special_tokens_mask(
+        self, token_ids_0: list[int], token_ids_1: Optional[list[int]] = None, already_has_special_tokens: bool = False
+    ) -> list[int]:
+        """
+        Retrieve 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` method.
+        Args:
+            token_ids_0 (`List[int]`):
+                List of IDs.
+            token_ids_1 (`List[int]`, *optional*):
+                Optional second list of IDs for sequence pairs.
+            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:
+            `List[int]`: A list of integers in the range [0, 1]: 1 for a special token, 0 for a sequence token.
+        """
+        if already_has_special_tokens:
+            return super().get_special_tokens_mask(
+                token_ids_0=token_ids_0, token_ids_1=token_ids_1, already_has_special_tokens=True
+            )
+
+        bos_token_id = [1] if self.add_bos_token else []
+        eos_token_id = [1] if self.add_eos_token else []
+
+        if token_ids_1 is None:
+            return bos_token_id + ([0] * len(token_ids_0)) + eos_token_id
+        return (
+            bos_token_id
+            + ([0] * len(token_ids_0))
+            + eos_token_id
+            + bos_token_id
+            + ([0] * len(token_ids_1))
+            + eos_token_id
+        )
+
+    # Copied from transformers.models.llama.tokenization_llama.LlamaTokenizer.create_token_type_ids_from_sequences
+    def create_token_type_ids_from_sequences(
+        self, token_ids_0: list[int], token_ids_1: Optional[list[int]] = None
+    ) -> list[int]:
+        """
+        Creates a mask from the two sequences passed to be used in a sequence-pair classification task. An ALBERT
+        sequence pair mask has the following format:
+        ```
+        0 0 0 0 0 0 0 0 0 0 0 1 1 1 1 1 1 1 1 1
+        | first sequence    | second sequence |
+        ```
+        if token_ids_1 is None, only returns the first portion of the mask (0s).
+        Args:
+            token_ids_0 (`List[int]`):
+                List of ids.
+            token_ids_1 (`List[int]`, *optional*):
+                Optional second list of IDs for sequence pairs.
+        Returns:
+            `List[int]`: List of [token type IDs](../glossary#token-type-ids) according to the given sequence(s).
+        """
+        bos_token_id = [self.bos_token_id] if self.add_bos_token else []
+        eos_token_id = [self.eos_token_id] if self.add_eos_token else []
+
+        output = [0] * len(bos_token_id + token_ids_0 + eos_token_id)
+
+        if token_ids_1 is not None:
+            output += [1] * len(bos_token_id + token_ids_1 + eos_token_id)
+
+        return output
+
+    def _tokenize(self, text: str, **kwargs) -> list[str]:
+        """Take as input a string and return a list of strings (tokens) for words/sub-words"""
+        tokens = self.convert_ids_to_tokens(self._bolmo_encode(text))
+        return tokens
+
+    def _patch_ids_to_byte_ids(self, input_ids: list[int]):
+        return [byte_token_id for token_id in input_ids for byte_token_id in self.byte_sequences[token_id]]
+
+    def _bolmo_encode(self, string: str, add_special_tokens=False):
+        input_ids = self.hf_tokenizer.encode(string, add_special_tokens=add_special_tokens)
+        return self._patch_ids_to_byte_ids(input_ids)
+
+    def _bolmo_decode(self, tokens: list[int], skip_special_tokens: bool = False) -> str:
+        return self._decode_to_bytes(tokens, skip_special_tokens=skip_special_tokens).decode("utf-8", errors="replace")
+
+    def _decode_to_bytes(self, tokens: list[int], skip_special_tokens: bool = False) -> bytes:
+        tokens_without_boundary = []
+        for token in tokens:
+            if token >= (self.offset + 256):
+                token -= self.offset + 256
+
+            tokens_without_boundary.append(token)
+
+        utf8_bytes = []
+
+        for token in tokens_without_boundary:
+            if token < self.offset:
+                if skip_special_tokens:
+                    continue
+                else:
+                    utf8_bytes.extend(self.config.special_tokens[token].encode("utf-8"))
+            else:
+                utf8_bytes.append(min(token - self.offset, 255))
+
+        return bytes(utf8_bytes)
+
+    def get_tokens_and_patch_lengths(self, original_input_ids: list[int], add_bos=False, strip_pad=False, skip_last=False):
+        if add_bos and self.bos_token_id is not None:
+            byte_tokens = [self.bos_token_id]
+            patch_lengths = [1]
+        else:
+            byte_tokens = []
+            patch_lengths = []
+
+        for idx, token in enumerate(original_input_ids):
+            # optionally skip last token to keep the length the same if add_bos=True
+            if skip_last and idx == len(original_input_ids) - 1:
+                break
+
+            token_byte_tokens = self._patch_ids_to_byte_ids([int(token)])
+
+            if strip_pad and all(t == self.pad_token_id for t in token_byte_tokens):
+                # skip padding tokens
+                continue
+
+            patch_lengths.append(len(token_byte_tokens))
+            byte_tokens.extend(token_byte_tokens)
+
+        return byte_tokens, patch_lengths
+
+    def convert_tokens_to_string(self, tokens: list[str]) -> str:
+        return self._bolmo_decode(self.convert_tokens_to_ids(tokens), skip_special_tokens=False)  # type: ignore
+
+    def _decode(
+        self,
+        token_ids: Union[int, list[int]],
+        skip_special_tokens: bool = False,
+        clean_up_tokenization_spaces: Optional[bool] = None,
+        spaces_between_special_tokens: bool = True,
+        **kwargs,
+    ) -> str:
+        if isinstance(token_ids, int):
+            token_ids = [token_ids]
+
+        return self._bolmo_decode(token_ids, skip_special_tokens=skip_special_tokens)
+
+    def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> tuple[str]:
+        return ()  # type: ignore

tokenizer_config.json

@@ -0,0 +1,34 @@
+{
+  "added_tokens_decoder": {
+    "0": {
+      "content": "<pad>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "<bos>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "auto_map": {
+    "AutoTokenizer": [
+      "tokenization_bolmo.BolmoTokenizer",
+      null
+    ]
+  },
+  "bos_token": "<bos>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<bos>",
+  "extra_ids": 0,
+  "extra_special_tokens": {},
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "<pad>",
+  "tokenizer_class": "BolmoTokenizer"
+}

utils_bolmo.py

@@ -0,0 +1,127 @@
+import math
+
+import torch
+import torch.nn.functional as F
+
+
+def compute_boundary_mask(boundary_logprobs: torch.Tensor, boundary_threshold: str) -> torch.Tensor:
+    if boundary_threshold.startswith("sample:"):
+        _, temperature = boundary_threshold.split(":")
+        temperature = float(temperature)
+
+        if temperature == 0:
+            return (boundary_logprobs > math.log(0.5))
+        elif temperature == 1:
+            return torch.bernoulli(torch.exp(boundary_logprobs)).to(torch.bool)
+        else:
+            raise NotImplementedError("Temperatures outside {0,1} are not implemented yet.")
+    elif boundary_threshold.startswith("topk:"):
+        _, topk = boundary_threshold.split(":")
+        topk = int(topk)
+        thresholds = torch.quantile(boundary_logprobs, dim=1, q=1 - (topk / boundary_logprobs.shape[1]))
+        return (boundary_logprobs >= thresholds.unsqueeze(-1))
+    elif boundary_threshold.startswith("topk_percent:"):
+        _, topk_percent = boundary_threshold.split(":")
+        topk_percent = float(topk_percent)
+        assert 0 <= topk_percent <= 1
+        thresholds = torch.quantile(boundary_logprobs, dim=1, q=1 - topk_percent)
+        return (boundary_logprobs >= thresholds.unsqueeze(-1))
+    else:
+        raise ValueError(f"Unknown boundary threshold: {boundary_threshold}")
+
+
+def _pad(tensors: list[torch.Tensor], multiple_of: int, direction: str, value):
+    max_len = max(t.size(0) for t in tensors)
+    if multiple_of > 1:
+        # Round up max_len to the nearest multiple_of
+        max_len = ((max_len + multiple_of - 1) // multiple_of) * multiple_of
+    padded = []
+    for t in tensors:
+        if direction == "left":
+            pad_shape = (max_len - t.size(0), 0)
+        elif direction == "right":
+            pad_shape = (0, max_len - t.size(0))
+        else:
+            raise ValueError(f"Unknown direction: {direction}. Must be 'left' or 'right'.")
+        padded.append(F.pad(t, pad_shape, value=value))
+    return torch.stack(padded, dim=0)
+
+def pad_right(
+    tensors: list[torch.Tensor],
+    multiple_of: int = 128,
+    value=0,
+):
+    return _pad(tensors, multiple_of, direction="right", value=value)
+
+def pad_left(
+    tensors: list[torch.Tensor],
+    multiple_of: int = 128,
+    value=0,
+):
+    return _pad(tensors, multiple_of, direction="left", value=value)
+
+class MaskState:
+    def __init__(self, mask):
+        self.cpu_mask = mask.cpu()
+
+        self.mask = mask
+        self.inv_mask = ~mask
+        self._all = self.cpu_mask.all().item()
+        self._any = self.cpu_mask.any().item()
+
+    def any(self):
+        return self._any
+
+    def all(self):
+        return self._all
+
+    def selective_get(self, x, inv=False):
+        # try to avoid sync through nonzero on index
+        if inv:
+            if self.all():
+                return x[[]]
+            elif not self.any():
+                return x
+            else:
+                return x[self.inv_mask]
+        else:
+            if self.all():
+                return x
+            elif not self.any():
+                return x[[]]
+            else:
+                return x[self.mask]
+
+    def selective_put(self, x, out, inv=False):
+        # try to avoid sync through nonzero on index
+        if inv:
+            if self.all():
+                return
+            elif not self.any():
+                out.copy_(x)
+            else:
+                out[self.inv_mask] = x
+        else:
+            if self.all():
+                out.copy_(x)
+            elif not self.any():
+                return
+            else:
+                out[self.mask] = x
+    
+    def selective_add(self, x, out, inv=False):
+        # try to avoid sync through nonzero on index
+        if inv:
+            if self.all():
+                return
+            elif not self.any():
+                out.add_(x)
+            else:
+                out[self.inv_mask] += x
+        else:
+            if self.all():
+                out.add_(x)
+            elif not self.any():
+                return
+            else:
+                out[self.mask] += x