lora finetune (need to be refactored)
This commit is contained in:
josc146
232
finetune/json2binidx_tool/tools/rwkv_tokenizer.py
vendored
Normal file
232
finetune/json2binidx_tool/tools/rwkv_tokenizer.py
vendored
Normal file
@@ -0,0 +1,232 @@
|
||||
########################################################################################################
|
||||
# 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 "<TRIE %s %s>"%(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()
|
||||
Reference in New Issue
Block a user