support for rwkv-4-world
This commit is contained in:
josc146
100525
backend-python/rwkv_pip/20B_tokenizer.json
Normal file
100525
backend-python/rwkv_pip/20B_tokenizer.json
Normal file
File diff suppressed because it is too large
Load Diff
106
backend-python/rwkv_pip/rwkv_tokenizer.py
Normal file
106
backend-python/rwkv_pip/rwkv_tokenizer.py
Normal file
@@ -0,0 +1,106 @@
|
||||
########################################################################################################
|
||||
# The RWKV Language Model - https://github.com/BlinkDL/RWKV-LM
|
||||
########################################################################################################
|
||||
|
||||
|
||||
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.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) -> list[int]:
|
||||
idx: int = 0
|
||||
tokens: list[int] = []
|
||||
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):
|
||||
try:
|
||||
return self.decodeBytes(tokens).decode("utf-8")
|
||||
except:
|
||||
return "\ufffd" # bad utf-8
|
||||
|
||||
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()
|
||||
65529
backend-python/rwkv_pip/rwkv_vocab_v20230424.txt
Normal file
65529
backend-python/rwkv_pip/rwkv_vocab_v20230424.txt
Normal file
File diff suppressed because it is too large
Load Diff
142
backend-python/rwkv_pip/utils.py
Normal file
142
backend-python/rwkv_pip/utils.py
Normal file
@@ -0,0 +1,142 @@
|
||||
########################################################################################################
|
||||
# The RWKV Language Model - https://github.com/BlinkDL/RWKV-LM
|
||||
########################################################################################################
|
||||
|
||||
import os, sys
|
||||
import numpy as np
|
||||
import torch
|
||||
from torch.nn import functional as F
|
||||
|
||||
|
||||
class PIPELINE_ARGS:
|
||||
def __init__(
|
||||
self,
|
||||
temperature=1.0,
|
||||
top_p=0.85,
|
||||
top_k=0,
|
||||
alpha_frequency=0.2,
|
||||
alpha_presence=0.2,
|
||||
token_ban=[],
|
||||
token_stop=[],
|
||||
chunk_len=256,
|
||||
):
|
||||
self.temperature = temperature
|
||||
self.top_p = top_p
|
||||
self.top_k = top_k
|
||||
self.alpha_frequency = alpha_frequency # Frequency Penalty (as in GPT-3)
|
||||
self.alpha_presence = alpha_presence # Presence Penalty (as in GPT-3)
|
||||
self.token_ban = token_ban # ban the generation of some tokens
|
||||
self.token_stop = token_stop # stop generation whenever you see any token here
|
||||
self.chunk_len = (
|
||||
chunk_len # split input into chunks to save VRAM (shorter -> slower)
|
||||
)
|
||||
|
||||
|
||||
class PIPELINE:
|
||||
def __init__(self, model, WORD_NAME):
|
||||
self.model = model
|
||||
if WORD_NAME == "cl100k_base":
|
||||
import tiktoken
|
||||
|
||||
self.tokenizer = tiktoken.get_encoding(WORD_NAME)
|
||||
elif WORD_NAME == "rwkv_vocab_v20230424":
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
from rwkv_tokenizer import TRIE_TOKENIZER
|
||||
|
||||
self.tokenizer = TRIE_TOKENIZER(
|
||||
os.path.dirname(os.path.abspath(__file__)) + "/rwkv_vocab_v20230424.txt"
|
||||
)
|
||||
else:
|
||||
from tokenizers import Tokenizer
|
||||
|
||||
self.tokenizer = Tokenizer.from_file(WORD_NAME)
|
||||
|
||||
def refine_context(self, context):
|
||||
context = context.strip().split("\n")
|
||||
for c in range(len(context)):
|
||||
context[c] = context[c].strip().strip("\u3000").strip("\r")
|
||||
context = list(filter(lambda c: c != "", context))
|
||||
context = "\n" + ("\n".join(context)).strip()
|
||||
if context == "":
|
||||
context = "\n"
|
||||
return context
|
||||
|
||||
def encode(self, x):
|
||||
if "Tokenizer" in str(type(self.tokenizer)):
|
||||
return self.tokenizer.encode(x).ids
|
||||
else:
|
||||
return self.tokenizer.encode(x)
|
||||
|
||||
def decode(self, x):
|
||||
return self.tokenizer.decode(x)
|
||||
|
||||
def sample_logits(self, logits, temperature=1.0, top_p=0.85, top_k=0):
|
||||
probs = F.softmax(logits.float(), dim=-1)
|
||||
top_k = int(top_k)
|
||||
if probs.device == torch.device("cpu"):
|
||||
probs = probs.numpy()
|
||||
sorted_ids = np.argsort(probs)
|
||||
sorted_probs = probs[sorted_ids][::-1]
|
||||
cumulative_probs = np.cumsum(sorted_probs)
|
||||
cutoff = float(sorted_probs[np.argmax(cumulative_probs > top_p)])
|
||||
probs[probs < cutoff] = 0
|
||||
if top_k < len(probs) and top_k > 0:
|
||||
probs[sorted_ids[:-top_k]] = 0
|
||||
if temperature != 1.0:
|
||||
probs = probs ** (1.0 / temperature)
|
||||
probs = probs / np.sum(probs)
|
||||
out = np.random.choice(a=len(probs), p=probs)
|
||||
return int(out)
|
||||
else:
|
||||
sorted_ids = torch.argsort(probs)
|
||||
sorted_probs = probs[sorted_ids]
|
||||
sorted_probs = torch.flip(sorted_probs, dims=(0,))
|
||||
cumulative_probs = torch.cumsum(sorted_probs, dim=-1).cpu().numpy()
|
||||
cutoff = float(sorted_probs[np.argmax(cumulative_probs > top_p)])
|
||||
probs[probs < cutoff] = 0
|
||||
if top_k < len(probs) and top_k > 0:
|
||||
probs[sorted_ids[:-top_k]] = 0
|
||||
if temperature != 1.0:
|
||||
probs = probs ** (1.0 / temperature)
|
||||
out = torch.multinomial(probs, num_samples=1)[0]
|
||||
return int(out)
|
||||
|
||||
def generate(
|
||||
self, ctx, token_count=100, args=PIPELINE_ARGS(), callback=None, state=None
|
||||
):
|
||||
all_tokens = []
|
||||
out_last = 0
|
||||
out_str = ""
|
||||
occurrence = {}
|
||||
for i in range(token_count):
|
||||
# forward & adjust prob.
|
||||
tokens = self.encode(ctx) if i == 0 else [token]
|
||||
while len(tokens) > 0:
|
||||
out, state = self.model.forward(tokens[: args.chunk_len], state)
|
||||
tokens = tokens[args.chunk_len :]
|
||||
|
||||
for n in args.token_ban:
|
||||
out[n] = -float("inf")
|
||||
for n in occurrence:
|
||||
out[n] -= args.alpha_presence + occurrence[n] * args.alpha_frequency
|
||||
|
||||
# sampler
|
||||
token = self.sample_logits(
|
||||
out, temperature=args.temperature, top_p=args.top_p, top_k=args.top_k
|
||||
)
|
||||
if token in args.token_stop:
|
||||
break
|
||||
all_tokens += [token]
|
||||
if token not in occurrence:
|
||||
occurrence[token] = 1
|
||||
else:
|
||||
occurrence[token] += 1
|
||||
|
||||
# output
|
||||
tmp = self.decode(all_tokens[out_last:])
|
||||
if "\ufffd" not in tmp: # is valid utf-8 string?
|
||||
if callback:
|
||||
callback(tmp)
|
||||
out_str += tmp
|
||||
out_last = i + 1
|
||||
return out_str
|
||||
Reference in New Issue
Block a user