110 lines
4.0 KiB
Python
110 lines
4.0 KiB
Python
import base64
|
|
import unicodedata
|
|
from typing import Collection, Dict, List, Set, Union
|
|
|
|
from .tokenizer_base import Tokenizer
|
|
|
|
PAT_STR = 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+""" # noqa E501
|
|
ENDOFTEXT = '<|endoftext|>'
|
|
IMSTART = '<|im_start|>'
|
|
IMEND = '<|im_end|>'
|
|
# as the default behavior is changed to allow special tokens in
|
|
# regular texts, the surface forms of special tokens need to be
|
|
# as different as possible to minimize the impact
|
|
EXTRAS = tuple((f'<|extra_{i}|>' for i in range(205)))
|
|
# changed to use actual index to avoid misconfiguration with vocabulary expansion
|
|
SPECIAL_START_ID = 151643
|
|
SPECIAL_TOKENS = tuple(
|
|
enumerate(
|
|
((
|
|
ENDOFTEXT,
|
|
IMSTART,
|
|
IMEND,
|
|
) + EXTRAS),
|
|
start=SPECIAL_START_ID,
|
|
))
|
|
SPECIAL_TOKENS_SET = set(t for i, t in SPECIAL_TOKENS)
|
|
|
|
|
|
class QwenTokenizer(Tokenizer):
|
|
@staticmethod
|
|
def _load_tiktoken_bpe(tiktoken_bpe_file: str) -> Dict[bytes, int]:
|
|
with open(tiktoken_bpe_file, 'rb') as f:
|
|
contents = f.read()
|
|
return {
|
|
base64.b64decode(token): int(rank)
|
|
for token, rank in (line.split() for line in contents.splitlines()
|
|
if line)
|
|
}
|
|
|
|
def __init__(self, vocab_file, errors='replace', extra_vocab_file=None):
|
|
self._errors = errors
|
|
self._vocab_file = vocab_file
|
|
self._extra_vocab_file = extra_vocab_file
|
|
|
|
self._mergeable_ranks = QwenTokenizer._load_tiktoken_bpe(
|
|
vocab_file) # type: Dict[bytes, int]
|
|
self._special_tokens = {
|
|
token: index
|
|
for index, token in SPECIAL_TOKENS
|
|
}
|
|
|
|
# try load extra vocab from file
|
|
if extra_vocab_file is not None:
|
|
used_ids = set(self._mergeable_ranks.values()) | set(
|
|
self._special_tokens.values())
|
|
extra_mergeable_ranks = self._load_tiktoken_bpe(extra_vocab_file)
|
|
for token, index in extra_mergeable_ranks.items():
|
|
if token in self._mergeable_ranks:
|
|
continue
|
|
if index in used_ids:
|
|
continue
|
|
self._mergeable_ranks[token] = index
|
|
# the index may be sparse after this, but don't worry tiktoken.Encoding will handle this
|
|
import tiktoken
|
|
enc = tiktoken.Encoding(
|
|
'Qwen',
|
|
pat_str=PAT_STR,
|
|
mergeable_ranks=self._mergeable_ranks,
|
|
special_tokens=self._special_tokens,
|
|
)
|
|
assert (
|
|
len(self._mergeable_ranks) +
|
|
len(self._special_tokens) == enc.n_vocab
|
|
), f'{len(self._mergeable_ranks) + len(self._special_tokens)} != {enc.n_vocab} in encoding'
|
|
|
|
self.decoder = {v: k
|
|
for k, v in self._mergeable_ranks.items()
|
|
} # type: dict[int, bytes|str]
|
|
self.decoder.update({v: k for k, v in self._special_tokens.items()})
|
|
|
|
self._tokenizer = enc # type: tiktoken.Encoding
|
|
|
|
self.eod_id = self._tokenizer.eot_token
|
|
self.im_start_id = self._special_tokens[IMSTART]
|
|
self.im_end_id = self._special_tokens[IMEND]
|
|
|
|
def encode(
|
|
self,
|
|
text: str,
|
|
allowed_special: Union[Set, str] = 'all',
|
|
disallowed_special: Union[Collection, str] = (),
|
|
) -> Union[List[List], List]:
|
|
text = unicodedata.normalize('NFC', text)
|
|
return self._tokenizer.encode(text,
|
|
allowed_special=allowed_special,
|
|
disallowed_special=disallowed_special)
|
|
|
|
def decode(
|
|
self,
|
|
token_ids: Union[int, List[int]],
|
|
skip_special_tokens: bool = False,
|
|
errors: str = None,
|
|
**kwargs,
|
|
) -> str:
|
|
if isinstance(token_ids, int):
|
|
token_ids = [token_ids]
|
|
if skip_special_tokens:
|
|
token_ids = [i for i in token_ids if i < self.eod_id]
|
|
return self._tokenizer.decode(token_ids, errors=errors or self._errors)
|