######################################################################################################## # The RWKV Language Model - https://github.com/BlinkDL/RWKV-LM # Source: https://github.com/BlinkDL/ChatRWKV/blob/main/tokenizer/rwkv_tokenizer.py ######################################################################################################## import os, sys, time, random print(''' ####################################################################################################################### This tokenizer is not used in any RWKV models yet. I plan to use it for the future multilang RWKV models. Benefits: * Good support of most languages, from European to CJK to Arabic and Hindi and more. * Clean vocab. Good for code too. Vocab size = 65525 (use 0 for <|endoftext|>). * Good at numbers: the numerical tokens are '0'~'9', '10'~'99', ' 0'~' 9', ' 10'~' 99'. * Very easy tokenization: ** The input text must be in UTF-8. ** Greedy encoding: always pick the longest (in bytes) token (with the highest id) that matches your UTF-8 bytes. * The tokenization result is surprisingly good, because the vocab respects word boundaries and UTF-8 boundaries. For 10x faster speed: mypyc rwkv_tokenizer.py python3 -c "import rwkv_tokenizer" ####################################################################################################################### ''') ######################################################################################################## # Tokenizer #1 (reference, naive, slow) ######################################################################################################## class RWKV_TOKENIZER(): table = None # : list[list[list[bytes]]] = None good = None # : list[set[int]] wlen = None # : list[int] def __init__(self, file_name): self.vocab_size = 65525 self.idx2token = {} sorted = [] # must be already sorted lines = open(file_name, "r", encoding="utf-8").readlines() for l in lines: idx = int(l[:l.index(' ')]) x = eval(l[l.index(' '):l.rindex(' ')]) x = x.encode("utf-8") if isinstance(x, str) else x assert isinstance(x, bytes) assert len(x) == int(l[l.rindex(' '):]) sorted += [x] self.idx2token[idx] = x self.token2idx = {} for k, v in self.idx2token.items(): self.token2idx[v] = int(k) # precompute some tables for fast matching self.table = [[[] for j in range(256)] for i in range(256)] self.good = [set() for i in range(256)] self.wlen = [0 for i in range(256)] for i in reversed(range(len(sorted))): # reverse order - match longer tokens first s = sorted[i] if len(s) >= 2: s0 = int(s[0]) s1 = int(s[1]) self.table[s0][s1] += [s] self.wlen[s0] = max(self.wlen[s0], len(s)) self.good[s0].add(s1) def encodeBytes(self, src: bytes): src_len: int = len(src) tokens = [] i: int = 0 while i < src_len: s: bytes = src[i : i + 1] if i < src_len - 1: s1: int = int(src[i + 1]) s0: int = int(src[i]) if s1 in self.good[s0]: sss: bytes = src[i : i + self.wlen[s0]] try: s = next(filter(sss.startswith, self.table[s0][s1])) except: pass tokens.append(self.token2idx[s]) i += len(s) return tokens def decodeBytes(self, tokens): return b''.join(map(lambda i: self.idx2token[i], tokens)) def encode(self, src: str): return self.encodeBytes(src.encode("utf-8")) def decode(self, tokens): return self.decodeBytes(tokens).decode('utf-8') def token_to_id(self, token): return self.token2idx[token] def get_vocab_size(self): return self.vocab_size def get_vocab(self): return self.idx2token def printTokens(self, tokens): for i in tokens: s = self.idx2token[i] try: s = s.decode('utf-8') except: pass print(f'{repr(s)}{i}', end=' ') # print(repr(s), i) print() ######################################################################################################## # Tokenizer #2 (trie, faster) https://github.com/TkskKurumi/ChatRWKV-TRIE-Tokenizer ######################################################################################################## class TRIE: __slots__ = tuple("ch,to,values,front".split(",")) to:list values:set def __init__(self, front=None, ch=None): self.ch = ch self.to = [None for ch in range(256)] self.values = set() self.front = front def __repr__(self): fr = self ret = [] while(fr!=None): if(fr.ch!=None): ret.append(fr.ch) fr = fr.front return ""%(ret[::-1], self.values) def add(self, key:bytes, idx:int=0, val=None): if(idx == len(key)): if(val is None): val = key self.values.add(val) return self ch = key[idx] if(self.to[ch] is None): self.to[ch] = TRIE(front=self, ch=ch) return self.to[ch].add(key, idx=idx+1, val=val) def find_longest(self, key:bytes, idx:int=0): u:TRIE = self ch:int = key[idx] while(u.to[ch] is not None): u = u.to[ch] idx += 1 if(u.values): ret = idx, u, u.values if(idx==len(key)): break ch = key[idx] return ret class TRIE_TOKENIZER(): def __init__(self, file_name): self.vocab_size = 65525 self.idx2token = {} sorted = [] # must be already sorted with open(file_name, "r", encoding="utf-8") as f: lines = f.readlines() for l in lines: idx = int(l[:l.index(' ')]) x = eval(l[l.index(' '):l.rindex(' ')]) x = x.encode("utf-8") if isinstance(x, str) else x assert isinstance(x, bytes) assert len(x) == int(l[l.rindex(' '):]) sorted += [x] self.idx2token[idx] = x self.token2idx = {} for k,v in self.idx2token.items(): self.token2idx[v] = int(k) self.root = TRIE() for t, i in self.token2idx.items(): _ = self.root.add(t, val=(t, i)) def encodeBytes(self, src:bytes): idx:int = 0 tokens = [] while (idx < len(src)): _idx:int = idx idx, _, values = self.root.find_longest(src, idx) assert(idx != _idx) _, token = next(iter(values)) tokens.append(token) return tokens def decodeBytes(self, tokens): return b''.join(map(lambda i: self.idx2token[i], tokens)) def encode(self, src): return self.encodeBytes(src.encode("utf-8")) def decode(self, tokens): return self.decodeBytes(tokens).decode('utf-8') def get_vocab_size(self): return self.vocab_size def get_vocab(self): return self.idx2token def printTokens(self, tokens): for i in tokens: s = self.idx2token[i] try: s = s.decode('utf-8') except: pass print(f'{repr(s)}{i}', end=' ') print()