| |
|
|
| 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}") |
|
|
| |
| 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") |
|
|
