# -*- coding:utf-8 -*- ######################################################################################################## # The RWKV v2-RNN Language Model - https://github.com/BlinkDL/RWKV-LM ######################################################################################################## import numpy as np import types import copy import torch from torch.nn import functional as F from src.utils import TOKENIZER torch.backends.cudnn.benchmark = True torch.backends.cudnn.allow_tf32 = True torch.backends.cuda.matmul.allow_tf32 = True ### Step 1: set model ################################################################################## ctx_len = 1024 n_layer = 6 n_embd = 512 model_type = 'RWKV' # 'RWKV' or 'RWKV-ffnPre' MODEL_NAME = 'trained-31' # your trained model WORD_NAME = 'vocab' # the .json vocab (generated by train.py) # --> set UNKNOWN_CHAR to the rarest token in your vocab.json <-- # --> unknown tokens in your context will be denoted by it <-- UNKNOWN_CHAR = ' ' # here we just set it to [space] for simplicity RUN_DEVICE = 'cpu' # 'cpu' (already very fast) or 'cuda' DEBUG_DEBUG = False # True False - show softmax output DEBUG_TIME = False # True False - show trained time-coeffs ### Step 2: set context ################################################################################ context = "\n" # ==> this is your prompt NUM_TRIALS = 999 LENGTH_PER_TRIAL = 500 TEMPERATURE = 1.0 top_p = 0.7 top_p_newline = 0.9 ######################################################################################################## np.set_printoptions(precision=4, suppress=True, linewidth=200) tokenizer = TOKENIZER(WORD_NAME, UNKNOWN_CHAR=UNKNOWN_CHAR) context = tokenizer.refine_context(context) print('Your context has ' + str(len(context)) + ' tokens') print(f'Loading {MODEL_NAME}...') ############################################################################################################## RWKV_K_CLAMP = 60 RWKV_K_EPS = 1e-16 RWKV_HEAD_QK_DIM = 256 class RWKV_RNN(): def __init__(self, MODEL_NAME): self.w = types.SimpleNamespace() w = torch.load(MODEL_NAME + '.pth', map_location=torch.device(RUN_DEVICE)) # .state_dict() for x in w.keys(): if '.time_' in x: w[x] = w[x].squeeze() if '.time_decay' in x: w[x] = torch.exp(-torch.exp(w[x])) if '.time_first' in x: w[x] = torch.exp(w[x]) xx = x.split('.') here = self.w for i in range(len(xx)): if xx[i].isdigit(): ii = int(xx[i]) if ii not in here: here[ii] = types.SimpleNamespace() here = here[ii] else: if i == len(xx) - 1: setattr(here, xx[i], w[x]) elif not hasattr(here, xx[i]): if xx[i+1].isdigit(): setattr(here, xx[i], {}) else: setattr(here, xx[i], types.SimpleNamespace()) here = getattr(here, xx[i]) self.clear() def clear(self): self.xx = {} self.aa = {} self.bb = {} self.hk = None def save(self, target): target.xx = copy.deepcopy(self.xx) target.aa = copy.deepcopy(self.aa) target.bb = copy.deepcopy(self.bb) target.hk = copy.deepcopy(self.hk) def load(self, target): self.xx = copy.deepcopy(target.xx) self.aa = copy.deepcopy(target.aa) self.bb = copy.deepcopy(target.bb) self.hk = copy.deepcopy(target.hk) def LN(self, xx, w): return F.layer_norm(xx, (n_embd,), weight=w.weight, bias=w.bias) def FF(self, xx, w, name): if DEBUG_TIME: print(name+'.time_mix', w.time_mix.squeeze().numpy()) if name not in self.xx: self.xx[name] = torch.zeros(n_embd, device=RUN_DEVICE) x = xx * w.time_mix + self.xx[name] * (1 - w.time_mix) self.xx[name] = xx r = torch.sigmoid(w.receptance.weight @ x) k = torch.square(torch.relu(w.key.weight @ x)) kv = w.value.weight @ k return r * kv def SA(self, xx, w, name): if DEBUG_TIME: print(name+'.time_mix', w.time_mix.squeeze().numpy()) print(name+'.time_decay', w.time_decay.squeeze().numpy()) print(name+'.time_first', w.time_first.squeeze().numpy()) if name not in self.xx: self.xx[name] = torch.zeros(n_embd, device=RUN_DEVICE) self.aa[name] = torch.zeros(n_embd, device=RUN_DEVICE) self.bb[name] = torch.zeros(n_embd, device=RUN_DEVICE) x = xx * w.time_mix + self.xx[name] * (1 - w.time_mix) self.xx[name] = xx r = torch.sigmoid(w.receptance.weight @ x) k = torch.exp(torch.clamp(w.key.weight @ x, max=RWKV_K_CLAMP)) v = w.value.weight @ x kv = k * v a = self.aa[name] + w.time_first * kv b = self.bb[name] + w.time_first * k self.aa[name] = w.time_decay * self.aa[name] + kv self.bb[name] = w.time_decay * self.bb[name] + k rwkv = r * a / (b + RWKV_K_EPS) return w.output.weight @ rwkv def run(self, ctx): w = self.w x = w.emb.weight[ctx[-1]] for i in range(n_layer): x = self.LN(x, w.blocks[i].ln1) if i == 0 and model_type == 'RWKV-ffnPre': x = x + self.FF(x, w.blocks[i].ffnPre, f'ffnPre.{i}') else: x = x + self.SA(x, w.blocks[i].att, f'att.{i}') x = self.LN(x, w.blocks[i].ln2) x = x + self.FF(x, w.blocks[i].ffn, f'ffn.{i}') x = self.LN(x, w.ln_out) if self.hk == None: self.hk = (w.head_k.weight @ x).unsqueeze(0) else: self.hk = torch.cat( [self.hk, (w.head_k.weight @ x).unsqueeze(0)], dim=0) if self.hk.shape[0] > ctx_len: self.hk = self.hk[-ctx_len:, :] q = w.head_q.weight @ x x = w.head.weight @ x x = x.cpu().numpy().tolist() c = (self.hk @ q) / RWKV_HEAD_QK_DIM for i in range(len(c)): x[ctx[i]] += c[i] return x ############################################################################################################## model = RWKV_RNN(MODEL_NAME) print('\n\n--> Currently the first run takes a while if your prompt is long, as we are using RNN to process the prompt. This will be much faster in future versions. <--\n') for TRIAL in range(1 if DEBUG_DEBUG else NUM_TRIALS): src_len = len(context) ctx = [tokenizer.stoi.get(s, tokenizer.UNKNOWN_CHAR) for s in context] print(context.replace('\n', '\n　'), end='') model.clear() if TRIAL == 0: init_state = types.SimpleNamespace() for i in range(src_len): x = ctx[:i+1] if i == src_len - 1: init_state.out = model.run(x) else: model.run(x) model.save(init_state) else: model.load(init_state) for i in range(src_len, src_len + (1 if DEBUG_DEBUG else LENGTH_PER_TRIAL)): x = ctx[:i+1] x = x[-ctx_len:] if i == src_len: out = copy.deepcopy(init_state.out) else: out = model.run(x) if DEBUG_DEBUG: print('model', np.array(x), '==>', np.array( out), np.max(out), np.min(out)) char = tokenizer.sample_logits(out, x, ctx_len, temperature=TEMPERATURE, top_p_usual=top_p, top_p_newline=top_p_newline) char = char.item() print(tokenizer.itos[int(char)].replace( '\n', '\n　'), end='', flush=True) ctx += [char] print('\n' + '-' * 40, end='')