import json import math import os from shutil import copyfile from typing import Any, Optional, Tuple import numpy as np # NOTE: numba does not support type hints for njit: https://github.com/python/mypy/issues/16149 from numba import njit # type: ignore[attr-defined] from numba.core import types from numba.typed import Dict, List from transformers.tokenization_utils import PreTrainedTokenizer from transformers.utils import logging VOCAB_FILES_NAMES = {"vocab_file": "tokenizer.jsonl"} logger = logging.get_logger(__name__) INVALID_SCORE = -20000000 UNKNOWN_SCORE = -10000000 TABLE_PIECE_LENGTH = 0 TABLE_TOKEN_ID = 1 TABLE_SCORE = 2 TABLE_PIECE_ID = 3 PATH_TOKEN_LENGTH = 0 PATH_TOKEN_ID = 1 PATH_NUM_TOKENS = 2 class AhoCorasick: def __init__(self) -> None: # List of tokens in the vocabulary. self._tokens: list[str] # A mapping from a byte code point to a token ID, used for byte fallback. self._bytes: np.ndarray # A mapping from a suffix's piece code to a suffix ID. # # Typically, the Aho-Corasick algorithm builds a Trie and adds suffix links between nodes # of the Trie. In this implementation, a suffix ID corresponds to a node in the trie, and # a piece code to an edge (in other words, a pair of a node and the next character). # # A piece code is a 64-bit integer: # - The upper 32 bits store the Unicode code point of the first character. # - The lower 32 bits store the suffix ID of the remaining suffix. # # A suffix ID is an integer indicating the starting position in the _table. self._to_suffix_id: Dict[types.int64, types.int32] # Flattened table representing the Trie structure for the Aho-Corasick algorithm. # It stores information including scores for each piece (prefix) within each suffix. # It is flattened for memory efficiency and performance. Suffixes are stored in # lexicographical order of their reversed strings, which improves memory access locality # when exploring new characters starting from the string's end. Pieces within a suffix are # stored in the decreasing order of their lengths. # # Each piece (a prefix fo the suffix) contains four pieces of information: # - TABLE_PIECE_LENGTH: Length of the piece. # - TABLE_TOKEN_ID: Token ID (or -1 if the piece is not a valid token). # - TABLE_SCORE: Score (or INVALID_SCORE if the piece is not a valid token). # - TABLE_PIECE_ID: Piece ID of the suffix. # # Each suffix also includes a sentinel row with a length of 1, a score of UNKNOWN_SCORE, # and a token ID of -1. Sentinel rows are identified by the score being UNKNOWN_SCORE. self._table: np.ndarray def build(self, vocab: list[Any]) -> None: self._bytes = np.zeros(256, dtype=np.int32) self._to_suffix_id = Dict.empty(key_type=types.int64, value_type=types.int32) # Build suffix_to_score and token_to_token_id. # The suffix_to_score dictionary maps a suffix to its score. It also includes all suffixes # of the token for the Trie structure for the Aho-Corasick algorithm. If a suffix is not a # valid token, its score is set to math.nan. # The token_to_token_id dictionary maps a token to its token ID. suffix_to_score: dict[str, float] = {} token_to_token_id: dict[str, int] = {} self._tokens = [] for token_id, row in enumerate(vocab): assert isinstance(row[0], str), row assert isinstance(row[1], (int, float)), row token = str(row[0]) self._tokens.append(token) token_to_token_id[token] = token_id # Special handling for byte tokens. if len(row) > 2 and row[2] == "BYTE": assert len(token) == 6 and token.startswith("<0x") and token.endswith(">"), row[0] self._bytes[int(row[0][3:5], 16)] = token_id continue suffix_to_score[token] = float(row[1]) # Ensure that all suffixes are included in suffix_to_score. for i in range(1, len(token)): suffix_to_score[token[i:]] = suffix_to_score.get(token[i:], math.nan) # Ensure all byte tokens are set. for i in range(256): assert self._bytes[i] != 0, f"Byte token for <0x{i:02X}> is not set." # List suffixes in lexicographical order of their reversed strings. suffixes = list(suffix_to_score.keys()) suffixes.append("") suffixes.sort(key=lambda x: x[::-1]) # Build suffix_to_id, which is a mapping from a suffix to a suffix ID, and _to_suffix_id, # which is a mapping from a piece code to a suffix ID. suffix_to_id: dict[str, int] = {} num_pieces = 0 for s in suffixes: suffix_to_id[s] = num_pieces if s != "": self._to_suffix_id[ord(s[0]) << 32 | suffix_to_id[s[1:]]] = np.int32(num_pieces) num_pieces += 1 + sum(s[:i] in suffix_to_score for i in range(1, len(s) + 1)) assert suffix_to_id[""] == 0, suffix_to_id[""] # Build _table, which is a flattened table representing the Trie structure for the Aho-Corasick. self._table = np.zeros((num_pieces, 4), dtype=np.int32) i = 0 for suffix in suffixes: # Add all prefixes of the suffix to the table. for piece_length in range(len(suffix), 0, -1): piece = suffix[:piece_length] score = suffix_to_score.get(piece, None) if score is None: continue self._table[i, TABLE_PIECE_LENGTH] = piece_length self._table[i, TABLE_TOKEN_ID] = token_to_token_id.get(piece, -1) self._table[i, TABLE_SCORE] = round(score * 1e4) if math.isfinite(score) else INVALID_SCORE self._table[i, TABLE_PIECE_ID] = suffix_to_id[piece] i += 1 # Add a sentinel row. self._table[i, TABLE_PIECE_LENGTH] = 1 self._table[i, TABLE_TOKEN_ID] = -1 self._table[i, TABLE_SCORE] = UNKNOWN_SCORE i += 1 assert i == num_pieces, (i, num_pieces) @staticmethod @njit def _encode( to_suffix_id: Dict[types.int64, types.int32], table: np.ndarray, bytes: np.ndarray, data: np.ndarray, ) -> np.ndarray: # Initialize scores array with a high value and set the score at the end to 0. # This array keeps track of the minimum cost (best score) to encode from each position to the end. scores = np.full((len(data) + 1,), 2**60, dtype=np.int64) scores[-1] = 0 # Path array to store the best path information. # The path array keeps track of token length, token ID, and number of tokens needed to encode. path = np.zeros((len(data) + 1, 3), dtype=np.int32) # Initialize suffix_id to 0, which represents the root of the Trie. suffix_id = 0 # Process the input data from the end to the beginning. for i in range(len(data) - 1, -1, -1): c = data[i] # Find the next suffix ID by iterating the suffix IDs of prefixes of the current suffix. # NOTE: If no suffix ID is found, suffix_id will be set to 0. for p in range(suffix_id, len(table)): suffix_id = to_suffix_id.get(c << 32 | table[p, TABLE_PIECE_ID], np.int32(0)) # If a next suffix ID is found or a sentinel row is reached, break the loop. if suffix_id > 0 or table[p, TABLE_SCORE] == UNKNOWN_SCORE: break # Update the best path to the current position. If multiple paths have the same score, # this chooses the longest prefix as the best path (table is sorted in the decreasing # order of piece length). for p in range(suffix_id, len(table)): score = table[p, TABLE_SCORE] if score > INVALID_SCORE: piece_length = table[p, TABLE_PIECE_LENGTH] s = scores[i + piece_length] - score if s < scores[i]: scores[i] = s path[i, PATH_TOKEN_LENGTH] = piece_length path[i, PATH_TOKEN_ID] = table[p, TABLE_TOKEN_ID] path[i, PATH_NUM_TOKENS] = path[i + piece_length, PATH_NUM_TOKENS] + 1 if score == UNKNOWN_SCORE: # Add number of bytes to represent `c` in UTF-8 (minus 1; 1 is already # added above). path[i, PATH_NUM_TOKENS] += (c >= 0x80) + (c >= 0x800) + (c >= 0x10000) # If it reaches a sentinel row, break the loop. if score == UNKNOWN_SCORE: break # Decode the best path from the beginning to get the token IDs. pos = 0 token_ids = np.zeros(path[0, PATH_NUM_TOKENS], dtype=np.int32) token_pos = 0 while pos < len(data): if path[pos, PATH_TOKEN_ID] >= 0: token_ids[token_pos] = path[pos, PATH_TOKEN_ID] token_pos += 1 else: # Fall back to byte tokens. c = data[pos] s = 1 + (c >= 0x80) + (c >= 0x800) + (c >= 0x10000) # Add byte tokens representing UTF-8 bytes. for i in range(s): b = c if s == 1 else (0xF00 >> s) & 0xFF if i == 0 else 0x80 token_ids[token_pos] = bytes[b | ((c >> (s - i - 1) * 6) & 0x3F)] token_pos += 1 # Ensure that pos should increase by at least 1. assert path[pos, PATH_TOKEN_LENGTH] > 0, (pos, path[pos]) pos += path[pos, PATH_TOKEN_LENGTH] return token_ids def encode(self, data: str) -> np.ndarray: """Encodes a string into a sequence of token IDs.""" return np.asarray( self._encode( self._to_suffix_id, self._table, self._bytes, # Convert a string into a numpy array of Unicode code points. # NOTE: This skips UTF-32 BOM. np.frombuffer(data.encode("utf-32"), dtype=np.int32)[1:], ) ) def encode_as_tokens(self, data: str) -> list[str]: """Encodes a string into a sequence of tokens.""" return [self._tokens[token_id] for token_id in self.encode(data)] class PlamoTokenizer(PreTrainedTokenizer): # type: ignore vocab_files_names = VOCAB_FILES_NAMES model_input_names = ["input_ids", "attention_mask"] _save_files = [ "special_tokens_map.json", "tokenization_plamo.py", "tokenizer.jsonl", "tokenizer_config.json", ] def __init__( self, vocab_file: str, unk_token: str = "<|plamo:unk|>", bos_token: str = "<|plamo:bos|>", eos_token: str = "<|plamo:eos|>", pad_token: str = "<|plamo:pad|>", cls_token: Optional[str] = None, sep_token: Optional[str] = None, mask_token: Optional[str] = None, clean_up_tokenization_spaces: bool = False, **kwargs: Any, ) -> None: """Tokenizer for PLaMo. Args: vocab_file (str): Vocabrary file path. unk_token (str): Unknown token. bos_token (str): Beginning of sentence token. eos_token (str): End of sentence token. pad_token (str): Padding token. cls_token (str): Classification token, to extract a summary of an input sequence leveraging self-attention along the full depth of the model. sep_token (str): Separation token, to separate context and query in an input sequence. mask_token (str): Mask token, to use when training a model with masked-language modeling. clean_up_tokenization_spaces (bool): Whether or not to clean up the tokenization spaces. num_threads (int): Number of threads. This value will be ignored if one of `PLAMO_TOKENIZER_NUM_THREADS` or `RAYON_NUM_THREADS` is set as an environment variable. """ if "add_bos_token" not in kwargs: kwargs["add_bos_token"] = False if "add_eos_token" not in kwargs: kwargs["add_eos_token"] = False self.data: list[Any] = [json.loads(line) for line in open(vocab_file, "r", encoding="utf-8")] self.vocab: dict[str, int] = {v[0]: i for i, v in enumerate(self.data)} self.aho_corasick = AhoCorasick() self.aho_corasick.build(self.data) self.vocab_file = vocab_file self.add_bos_token = kwargs["add_bos_token"] self.add_eos_token = kwargs["add_eos_token"] super().__init__( vocab_file=vocab_file, unk_token=unk_token, bos_token=bos_token, eos_token=eos_token, pad_token=pad_token, cls_token=cls_token, sep_token=sep_token, mask_token=mask_token, clean_up_tokenization_spaces=clean_up_tokenization_spaces, **kwargs, ) # the functions below are copied from hf transformers LlamaTokenizer's implementation to fix the behaviour of the tokenizer # https://github.com/huggingface/transformers/blob/v4.30.2/src/transformers/models/llama/tokenization_llama.py def __getstate__(self) -> dict[str, Any]: state = self.__dict__.copy() state["aho_corasick"] = None return state def __setstate__(self, d: dict[str, Any]) -> None: self.__dict__ = d self.aho_corasick = AhoCorasick() self.aho_corasick.build(self.data) @property def vocab_size(self) -> Any: """Returns vocab size""" return len(self.data) def token_to_score(self, token: str) -> Optional[float]: """Returns score of the token""" token_id = self.vocab.get(token, None) return None if token_id is None else self.data[token_id][1] def get_vocab(self) -> dict[str, int]: """Returns vocab as a dict""" vocab = self.vocab.copy() vocab.update(self.added_tokens_encoder) return vocab def convert_tokens_to_string(self, tokens: List[str]) -> str: """Converts a sequence of tokens (string) in a single string.""" return b"".join( [bytes([int(t[3:5], 16)]) if t.startswith("<0x") else t.encode("utf-8") for t in tokens] ).decode("utf-8", errors="replace") def _tokenize(self, text: str) -> Any: """Returns a tokenized string.""" return self.aho_corasick.encode_as_tokens(text) def _convert_token_to_id(self, token: str) -> Any: """Converts a token (str) in an id using the vocab.""" return self.vocab.get(token, 0) def _convert_id_to_token(self, index: int) -> Any: """Converts an index (integer) in a token (str) using the vocab.""" return self.data[index][0] def build_inputs_with_special_tokens( self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None ) -> List[int]: 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 def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]: """ Save the vocabulary and special tokens file to a directory. Args: save_directory (`str`): The directory in which to save the vocabulary. Returns: `Tuple(str)`: Paths to the files saved. """ if not os.path.isdir(save_directory): logger.error(f"Vocabulary path ({save_directory}) should be a directory") return ("",) out_vocab_file = os.path.join( save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"] ) if os.path.abspath(self.vocab_file) != os.path.abspath(out_vocab_file) and os.path.isfile(self.vocab_file): copyfile(self.vocab_file, out_vocab_file) elif not os.path.isfile(self.vocab_file): with open(out_vocab_file, "w") as f: for token in self.data: print(json.dumps(token, ensure_ascii=False), file=f) return (out_vocab_file,)