| |
| |
| import base64 |
| import shutil |
| import unicodedata |
| from pathlib import Path |
|
|
| import tiktoken |
| from transformers import PreTrainedTokenizer |
|
|
|
|
| class YuE2Tokenizer(PreTrainedTokenizer): |
| vocab_files_names = {"vocab_file": "qwen.tiktoken"} |
| model_input_names = ["input_ids"] |
|
|
| def __init__(self, vocab_file, expected_vocab_size=151643, **kwargs): |
| self.vocab_file = str(vocab_file) |
| self.ranks = { |
| base64.b64decode(token): int(rank) |
| for token, rank in (line.split() for line in Path(vocab_file).read_bytes().splitlines() if line) |
| } |
| if len(self.ranks) != expected_vocab_size: |
| raise ValueError(f"Expected {expected_vocab_size} ordinary tokens") |
| specials = ["<|endoftext|>", "<|im_start|>", "<|im_end|>", "<R>", "<S>", "<X>", "<mask>", "<sep>"] |
| specials += [f"<extra_{i}>" for i in range(200)] |
| specials[204:206] = ["<abc>", "</abc>"] |
| pattern = 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+" |
| self.encoding = tiktoken.Encoding( |
| "YuE2", |
| pat_str=pattern, |
| mergeable_ranks=self.ranks, |
| special_tokens={s: i + len(self.ranks) for i, s in enumerate(specials)}, |
| ) |
| super().__init__(expected_vocab_size=expected_vocab_size, **kwargs) |
|
|
| @property |
| def vocab_size(self): |
| return len(self.ranks) |
|
|
| def get_vocab(self): |
| return {base64.b64encode(token).decode(): rank for token, rank in self.ranks.items()} |
|
|
| def _tokenize(self, text): |
| return [base64.b64encode(self.encoding.decode_single_token_bytes(i)).decode() for i in self.encode(text)] |
|
|
| def _convert_token_to_id(self, token): |
| return self.ranks[base64.b64decode(token)] |
|
|
| def _convert_id_to_token(self, index): |
| return base64.b64encode(self.encoding.decode_single_token_bytes(index)).decode() |
|
|
| def encode(self, text, **kwargs): |
| return self.encoding.encode_ordinary(unicodedata.normalize("NFC", text)) |
|
|
| def decode(self, token_ids, **kwargs): |
| return self.encoding.decode([int(i) for i in token_ids if 0 <= i < self.encoding.n_vocab], errors="replace") |
|
|
| def save_vocabulary(self, save_directory, filename_prefix=None): |
| target = Path(save_directory) / ((filename_prefix + "-" if filename_prefix else "") + "qwen.tiktoken") |
| if target.resolve() != Path(self.vocab_file).resolve(): |
| shutil.copyfile(self.vocab_file, target) |
| return (str(target),) |
|
|
