# Copyright (c) Meta Platforms, Inc. and affiliates. import abc import logging import os from copy import copy from dataclasses import dataclass from pathlib import Path from typing import List, Optional, Tuple import tiktoken from sentencepiece import SentencePieceProcessor from tiktoken.load import load_tiktoken_bpe logger = logging.getLogger(__name__) @dataclass class TokenizerArgs: name: str = "bytes" path: Optional[str] = None class Tokenizer(abc.ABC): @abc.abstractmethod def encode(self, tokens, add_bos, add_eos): pass @abc.abstractmethod def decode(self, tokens): pass @abc.abstractmethod def get_token_offsets( self, text: str, tokens: Optional[List[int]] = None ) -> Tuple[List[str], List[int]]: """Return the offsets of the tokens in the original text. Only used for evaluation.""" pass class MockTokenizer(Tokenizer): n_words: int = 256 def encode(self, tokens, add_bos, add_eos): return tokens class ByteTokenizer(Tokenizer): def __init__(self): self.bos_id = 256 self.eos_id = 257 self.n_words = 258 def encode(self, s: str, add_bos: bool = False, add_eos: bool = False): tokens = [self.bos_id] * add_bos + list(s.encode()) + [self.eos_id] * add_eos return tokens def decode(self, tokens: List[int]): byte_tokens = bytes([t for t in tokens if t < 256]) return byte_tokens.decode("utf-8", errors="backslashreplace") def get_token_offsets( self, text: str, tokens: Optional[List[int]] = None ) -> Tuple[List[str], List[int]]: if tokens is None: tokens = self.encode(text) decoded_chars, offsets = [], [] byte_pos = 0 for token in tokens: if token < 256: char = bytes([token]).decode("utf-8", errors="ignore") if char: decoded_chars.append(char) offsets.append(byte_pos) byte_pos += len(char.encode("utf-8")) return decoded_chars, offsets class SentencePieceTokenizer(Tokenizer): def __init__(self, model_path: str) -> None: assert os.path.isfile(model_path), model_path self.sp_model = SentencePieceProcessor(model_file=model_path) logger.info(f"Reloaded SentencePiece model from {model_path}") # BOS / EOS token IDs self.n_words: int = self.sp_model.vocab_size() self.bos_id: int = self.sp_model.bos_id() self.eos_id: int = self.sp_model.eos_id() self.pad_id: int = self.sp_model.pad_id() logger.info( f"#words: {self.n_words} - BOS ID: {self.bos_id} - EOS ID: {self.eos_id}" ) assert self.sp_model.vocab_size() == self.sp_model.get_piece_size() def encode(self, s: str, add_bos: bool, add_eos: bool): assert type(s) is str tokens = ( [self.bos_id] * add_bos + self.sp_model.encode(s) + [self.eos_id] * add_eos ) return tokens def decode(self, tokens: List[int]): return self.sp_model.decode(tokens) def get_token_offsets( self, text: str, tokens: Optional[List[int]] = None ) -> Tuple[List[str], List[int]]: pieces = self.sp_model.encode_as_immutable_proto(text).pieces substrs = [p.surface for p in pieces] offsets = [p.begin for p in pieces] return substrs, offsets DEFAULT_TIKTOKEN_PATTERN = r"""(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+""" DEFAULT_TIKTOKEN_SPECIAL_TOKENS = { "<|begin_of_text|>": 0, "<|end_of_text|>": 1, "<|fim_prefix|>": 2, "<|fim_middle|>": 3, "<|fim_end_fill|>": 253, "<|fim_pad|>": 254, "<|fim_suffix|>": 255, } TIKTOKEN_MAX_ENCODE_CHARS = 400_000 class TikTokenTokenizer(Tokenizer): def __init__(self, model_path: str) -> None: mergeable_ranks = load_tiktoken_bpe(model_path) all_special_tokens_with_ids = copy(DEFAULT_TIKTOKEN_SPECIAL_TOKENS) missing_ids = set(range(256)) - set(all_special_tokens_with_ids.values()) for id in missing_ids: all_special_tokens_with_ids[f"<|reserved_special_token_{id}|>"] = id for name in all_special_tokens_with_ids: all_special_tokens_with_ids[name] += len(mergeable_ranks) self.tkt_model = tiktoken.core.Encoding( name=Path(model_path).stem, pat_str=DEFAULT_TIKTOKEN_PATTERN, mergeable_ranks=mergeable_ranks, special_tokens=all_special_tokens_with_ids, ) self.bos_id: int = self.tkt_model.encode_single_token("<|begin_of_text|>") self.eos_id: int = self.tkt_model.encode_single_token("<|end_of_text|>") self.n_words: int = self.tkt_model.n_vocab logger.info( f"#words: {self.n_words} - BOS ID: {self.bos_id} - EOS ID: {self.eos_id}" ) def encode(self, s: str, add_bos: bool, add_eos: bool): assert isinstance(s, str) subs = [] for i in range(0, len(s), TIKTOKEN_MAX_ENCODE_CHARS): subs.append(s[i : i + TIKTOKEN_MAX_ENCODE_CHARS]) return ( [self.bos_id] * add_bos + sum(self.tkt_model.encode_ordinary_batch(subs), start=[]) + [self.eos_id] * add_eos ) def decode(self, tokens: List[int]): return self.tkt_model.decode(tokens) def get_token_offsets( self, text: str, tokens: Optional[List[int]] = None ) -> Tuple[List[str], List[int]]: if tokens is not None: token_bytes = self.tkt_model.decode_tokens_bytes(tokens) else: token_bytes = self.tkt_model.decode_tokens_bytes( self.tkt_model.encode(text, allowed_special="all") ) text_len, offsets = 0, [] for token in token_bytes: offsets.append(max(0, text_len - (0x80 <= token[0] < 0xC0))) text_len += sum(1 for c in token if not 0x80 <= c < 0xC0) substrs = [text[s:e] for s, e in zip(offsets, offsets[1:] + [None])] return substrs, offsets def build_tokenizer(name: str, path: Optional[str] = None) -> Tokenizer: if name == "bytes": return ByteTokenizer() elif name == "mock": return MockTokenizer() elif name == "sp": return SentencePieceTokenizer(path) elif name == "tiktoken": return TikTokenTokenizer(path) else: raise NotImplementedError(f"{name} tokenizer type is not implemented")