Upload 6 files

MiniMind2_tokenizer/special_tokens_map.json

@@ -0,0 +1,30 @@
+{
+  "bos_token": {
+    "content": "<|im_start|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "<|im_end|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

MiniMind2_tokenizer/tokenizer_config.json

@@ -0,0 +1,44 @@
+{
+  "add_bos_token": false,
+  "add_eos_token": false,
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "0": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "<|im_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "<|im_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "additional_special_tokens": [],
+  "bos_token": "<|im_start|>",
+  "chat_template": "{% if messages[0]['role'] == 'system' %}{% set system_message = messages[0]['content'] %}{{ '<|im_start|>system\\n' + system_message + '<|im_end|>\\n' }}{% else %}{{ '<|im_start|>system\\nYou are a helpful assistant<|im_end|>\\n' }}{% endif %}{% for message in messages %}{% set content = message['content'] %}{% if message['role'] == 'user' %}{{ '<|im_start|>user\\n' + content + '<|im_end|>\\n<|im_start|>assistant\\n' }}{% elif message['role'] == 'assistant' %}{{ content + '<|im_end|>' + '\\n' }}{% endif %}{% endfor %}",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "extra_special_tokens": {},
+  "legacy": true,
+  "model_max_length": 32768,
+  "pad_token": "<|endoftext|>",
+  "sp_model_kwargs": {},
+  "spaces_between_special_tokens": false,
+  "tokenizer_class": "PreTrainedTokenizer",
+  "unk_token": "<|endoftext|>"
+}

app.py

@@ -0,0 +1,402 @@
+import types, torch, copy
+from typing import List
+torch._C._jit_set_autocast_mode(False)
+import torch.nn as nn
+from torch.nn import functional as F
+from transformers import AutoTokenizer
+import gradio as gr
+MyModule = torch.jit.ScriptModule
+MyFunction = torch.jit.script_method
+MyStatic = torch.jit.script
+
+########################################################################################################
+
+args = types.SimpleNamespace()
+args.MODEL_NAME = "./sft-2048.pth"
+args.n_layer = 8
+args.n_embd = 512
+args.vocab_size = 6400
+args.head_size = 64
+GEN_TEMP = 1.0
+GEN_TOP_P = 0.3
+GEN_alpha_presence = 0.5
+GEN_alpha_frequency = 0.5
+GEN_penalty_decay = 0.996
+CHUNK_LEN = 16
+DTYPE = torch.float32
+HEAD_SIZE = args.head_size
+STATE_NAME = None
+
+########################################################################################################
+
+class RWKV_x070(MyModule):
+    def __init__(self, args):
+        super().__init__()
+        self.args = args
+        self.n_embd = args.n_embd
+        self.n_layer = args.n_layer
+        self.eval()
+        
+        self.z = torch.load(args.MODEL_NAME, map_location='cuda')
+        z = self.z
+        self.n_head, self.head_size = z['blocks.0.att.r_k'].shape
+
+        keys = list(z.keys())
+        for k in keys:
+            if 'key.weight' in k or 'value.weight' in k or 'receptance.weight' in k or 'output.weight' in k or 'head.weight' in k:
+                z[k] = z[k].t()
+            z[k] = z[k].squeeze().to(dtype=DTYPE)
+            if k.endswith('att.r_k'): z[k] = z[k].flatten()
+        assert self.head_size == args.head_size
+
+        z['emb.weight'] = F.layer_norm(z['emb.weight'], (args.n_embd,), weight=z['blocks.0.ln0.weight'], bias=z['blocks.0.ln0.bias'])
+
+        for i in range(self.n_layer): # !!! merge emb residual !!!
+            z[f'blocks.{i}.ffn.s_emb.weight'] = z[f'blocks.{i}.ffn.s_emb.weight'] + z['emb.weight'] @ z[f'blocks.{i}.ffn.s_emb_x.weight'].t()
+
+        z['blocks.0.att.v0'] = z['blocks.0.att.a0'] # actually ignored
+        z['blocks.0.att.v1'] = z['blocks.0.att.a1'] # actually ignored
+        z['blocks.0.att.v2'] = z['blocks.0.att.a2'] # actually ignored
+
+    def forward(self, idx, state, full_output=False):
+        if state == None:
+            state = [None for _ in range(args.n_layer * 3)]
+            for i in range(args.n_layer): # state: 0=att_x_prev 1=att_kv 2=ffn_x_prev
+                state[i*3+0] = torch.zeros(args.n_embd, dtype=DTYPE, requires_grad=False, device="cuda")
+                state[i*3+1] = torch.zeros((args.n_embd // args.head_size, args.head_size, args.head_size), dtype=torch.float, requires_grad=False, device="cuda")
+                state[i*3+2] = torch.zeros(args.n_embd, dtype=DTYPE, requires_grad=False, device="cuda")
+
+        if type(idx) is list:
+            if len(idx) > 1:
+                return self.forward_seq(idx, state, full_output)
+            else:
+                return self.forward_one(idx[0], state)
+        else:
+            return self.forward_one(idx, state)
+
+    @MyFunction
+    def forward_one(self, idx:int, state:List[torch.Tensor]):
+        with torch.no_grad(): 
+            z = self.z
+            x = z['emb.weight'][idx]
+
+            v_first = torch.empty_like(x)
+            for i in range(self.n_layer):
+                bbb = f'blocks.{i}.'
+                att = f'blocks.{i}.att.'
+                ffn = f'blocks.{i}.ffn.'
+
+                xx = F.layer_norm(x, (self.n_embd,), weight=z[bbb+'ln1.weight'], bias=z[bbb+'ln1.bias'])
+
+                xx, state[i*3+0], state[i*3+1], v_first = RWKV_x070_TMix_one(i, self.n_head, self.head_size, xx, state[i*3+0], v_first, state[i*3+1],
+                    z[att+'x_r'], z[att+'x_w'], z[att+'x_k'], z[att+'x_v'], z[att+'x_a'], z[att+'x_g'],
+                    z[att+'w0'], z[att+'w1'], z[att+'w2'], z[att+'a0'], z[att+'a1'], z[att+'a2'], z[att+'v0'], z[att+'v1'], z[att+'v2'],
+                    z[att+'g1'], z[att+'g2'], z[att+'k_k'], z[att+'k_a'], z[att+'r_k'],
+                    z[att+'receptance.weight'], z[att+'key.weight'], z[att+'value.weight'], z[att+'output.weight'],
+                    z[att+'ln_x.weight'], z[att+'ln_x.bias'])
+                x = x + xx
+
+                xx = F.layer_norm(x, (self.n_embd,), weight=z[bbb+'ln2.weight'], bias=z[bbb+'ln2.bias'])
+
+                xx, state[i*3+2] = RWKV_x070_CMix_one(xx, state[i*3+2], z[ffn+'x_k'], z[ffn+'key.weight'], z[ffn+'value.weight'], z[ffn+'s_emb.weight'][idx], z[ffn+'s1'], z[ffn+'s2'], z[ffn+'s0'])
+                x = x + xx
+            
+            x = F.layer_norm(x, (self.n_embd,), weight=z['ln_out.weight'], bias=z['ln_out.bias'])
+            x = x @ z['head.weight']
+            return x, state
+        
+    @MyFunction
+    def forward_seq(self, idx:List[int], state:List[torch.Tensor], full_output:bool=False):
+        with torch.no_grad(): 
+            z = self.z
+            x = z['emb.weight'][idx]
+
+            v_first = torch.empty_like(x)
+            for i in range(self.n_layer):
+                bbb = f'blocks.{i}.'
+                att = f'blocks.{i}.att.'
+                ffn = f'blocks.{i}.ffn.'
+
+                xx = F.layer_norm(x, (self.n_embd,), weight=z[bbb+'ln1.weight'], bias=z[bbb+'ln1.bias'])
+
+                xx, state[i*3+0], state[i*3+1], v_first = RWKV_x070_TMix_seq(i, self.n_head, self.head_size, xx, state[i*3+0], v_first, state[i*3+1],
+                    z[att+'x_r'], z[att+'x_w'], z[att+'x_k'], z[att+'x_v'], z[att+'x_a'], z[att+'x_g'],
+                    z[att+'w0'], z[att+'w1'], z[att+'w2'], z[att+'a0'], z[att+'a1'], z[att+'a2'], z[att+'v0'], z[att+'v1'], z[att+'v2'],
+                    z[att+'g1'], z[att+'g2'], z[att+'k_k'], z[att+'k_a'], z[att+'r_k'],
+                    z[att+'receptance.weight'], z[att+'key.weight'], z[att+'value.weight'], z[att+'output.weight'],
+                    z[att+'ln_x.weight'], z[att+'ln_x.bias'])
+                x = x + xx
+
+                xx = F.layer_norm(x, (self.n_embd,), weight=z[bbb+'ln2.weight'], bias=z[bbb+'ln2.bias'])
+
+                xx, state[i*3+2] = RWKV_x070_CMix_seq(xx, state[i*3+2], z[ffn+'x_k'], z[ffn+'key.weight'], z[ffn+'value.weight'], z[ffn+'s_emb.weight'][idx], z[ffn+'s1'], z[ffn+'s2'], z[ffn+'s0'])
+                x = x + xx
+            
+            if not full_output: x = x[-1,:]
+            x = F.layer_norm(x, (self.n_embd,), weight=z['ln_out.weight'], bias=z['ln_out.bias'])
+            x = x @ z['head.weight']
+            return x, state
+
+########################################################################################################
+
+@MyStatic
+def RWKV_x070_TMix_one(layer_id: int, H:int, N:int, x, x_prev, v_first, state, x_r, x_w, x_k, x_v, x_a, x_g, w0, w1, w2, a0, a1, a2, v0, v1, v2, g1, g2, k_k, k_a, r_k, R_, K_, V_, O_, ln_w, ln_b):
+    xx = x_prev - x
+    xr, xw, xk, xv, xa, xg = x+xx*x_r, x+xx*x_w, x+xx*x_k, x+xx*x_v, x+xx*x_a, x+xx*x_g
+
+    r = xr @ R_
+    w = torch.tanh(xw @ w1) @ w2
+    k = xk @ K_
+    v = xv @ V_
+    a = torch.sigmoid(a0 + (xa @ a1) @ a2)
+    g = torch.sigmoid(xg @ g1) @ g2
+
+    kk = torch.nn.functional.normalize((k * k_k).view(H,N), dim=-1, p=2.0).view(H*N)
+    k = k * (1 + (a-1) * k_a)
+    if layer_id == 0: v_first = v
+    else: v = v + (v_first - v) * torch.sigmoid(v0 + (xv @ v1) @ v2)
+    w = torch.exp(-0.606531 * torch.sigmoid((w0 + w).float())) # 0.606531 = exp(-0.5)
+
+    vk = v.view(H,N,1) @ k.view(H,1,N)
+    ab = (-kk).view(H,N,1) @ (kk*a).view(H,1,N)
+    state = state * w.view(H,1,N) + state @ ab.float() + vk.float()
+    xx = (state.to(dtype=x.dtype) @ r.view(H,N,1))
+
+    xx = torch.nn.functional.group_norm(xx.view(1,H*N), num_groups=H, weight=ln_w, bias=ln_b, eps = 64e-5).view(H*N)    
+    xx = xx + ((r * k * r_k).view(H,N).sum(dim=-1, keepdim=True) * v.view(H,N)).view(H*N)
+    return (xx * g) @ O_, x, state, v_first
+
+@MyStatic
+def RWKV_x070_TMix_seq(layer_id: int, H:int, N:int, x, x_prev, v_first, state, x_r, x_w, x_k, x_v, x_a, x_g, w0, w1, w2, a0, a1, a2, v0, v1, v2, g1, g2, k_k, k_a, r_k, R_, K_, V_, O_, ln_w, ln_b):
+    T = x.shape[0]
+    xx = torch.cat((x_prev.unsqueeze(0), x[:-1,:])) - x
+    xr, xw, xk, xv, xa, xg = x+xx*x_r, x+xx*x_w, x+xx*x_k, x+xx*x_v, x+xx*x_a, x+xx*x_g
+
+    r = xr @ R_
+    w = torch.tanh(xw @ w1) @ w2
+    k = xk @ K_
+    v = xv @ V_
+    a = torch.sigmoid(a0 + (xa @ a1) @ a2)
+    g = torch.sigmoid(xg @ g1) @ g2
+
+    kk = torch.nn.functional.normalize((k * k_k).view(T,H,N), dim=-1, p=2.0).view(T,H*N)
+    k = k * (1 + (a-1) * k_a)
+    if layer_id == 0: v_first = v
+    else: v = v + (v_first - v) * torch.sigmoid(v0 + (xv @ v1) @ v2)
+
+    ######## cuda-free method 
+    w = torch.exp(-0.606531 * torch.sigmoid((w0 + w).float())) # 0.606531 = exp(-0.5)
+    for t in range(T):
+        r_, w_, k_, v_, kk_, a_ = r[t], w[t], k[t], v[t], kk[t], a[t]
+        vk = v_.view(H,N,1) @ k_.view(H,1,N)
+        ab = (-kk_).view(H,N,1) @ (kk_*a_).view(H,1,N)
+        state = state * w_.view(H,1,N) + state @ ab.float() + vk.float()
+        xx[t] = (state.to(dtype=x.dtype) @ r_.view(H,N,1)).view(H*N)
+
+    # w = -torch.nn.functional.softplus(-(w0 + w)) - 0.5
+    # xx = RWKV7_OP(state, r, w, k, v, -kk, kk*a)
+
+    xx = torch.nn.functional.group_norm(xx.view(T,H*N), num_groups=H, weight=ln_w, bias=ln_b, eps = 64e-5).view(T,H*N)
+    xx = xx + ((r * k * r_k).view(T,H,N).sum(dim=-1, keepdim=True) * v.view(T,H,N)).view(T,H*N)
+    return (xx * g) @ O_, x[-1,:], state, v_first
+
+########################################################################################################
+
+@MyStatic
+def RWKV_x070_CMix_one(x, x_prev, x_k, K_, V_, semb_, s1_, s2_, s0_):
+    xx = x_prev - x
+    k = x + xx * x_k
+    k = torch.relu(k @ K_) ** 2
+    ss = (x @ s1_) @ semb_.view(32,32)
+    k = k * ((ss @ s2_) + s0_)
+    return k @ V_, x
+
+@MyStatic
+def RWKV_x070_CMix_seq(x, x_prev, x_k, K_, V_, semb_, s1_, s2_, s0_):
+    T,C = x.shape
+    xx = torch.cat((x_prev.unsqueeze(0), x[:-1,:])) - x
+    k = x + xx * x_k
+    k = torch.relu(k @ K_) ** 2    
+    ss = (x @ s1_).view(T,1,32) @ semb_.view(T,32,32)
+    k = k * ((ss.view(T,32) @ s2_) + s0_)
+    return k @ V_, x[-1,:]
+
+@MyStatic
+def sample_logits(logits, temperature:float=1.0, top_p:float=1.0, top_k:int=0):
+    probs = F.softmax(logits.float(), dim=-1)
+    sorted_probs, sorted_ids = torch.sort(probs, descending=True)
+    
+    if top_k > 0:
+        probs[sorted_ids[top_k:]] = 0
+
+    if top_p < 1:
+        cumulative_probs = torch.cumsum(sorted_probs, dim=-1)
+        cutoff_index = torch.searchsorted(cumulative_probs, top_p)
+        cutoff = sorted_probs[cutoff_index]
+        probs[probs < cutoff] = 0
+
+        if top_p > 0:
+            idx = torch.where(probs == cutoff)[0]
+            if len(idx) > 0:
+                probs[idx] = cutoff + (top_p - torch.sum(probs).item()) / len(idx)
+                # assert abs(torch.sum(probs).item() - top_p) < 1e-6
+    
+    if temperature != 1.0:
+        probs = probs ** (1.0 / temperature)
+
+    return torch.multinomial(probs, num_samples=1).item()
+
+tokenizer = AutoTokenizer.from_pretrained("./MiniMind2_tokenizer")
+
+model_tokens = []
+model_state = None
+model = RWKV_x070(args)
+
+# if STATE_NAME is not None:
+#     GEN_TOP_P = 0.2
+#     GEN_alpha_presence = 0.3
+#     GEN_alpha_frequency = 0.3
+    
+#     args = model.args
+#     state_raw = torch.load(STATE_NAME + '.pth')
+#     state_init = [None for i in range(args.n_layer * 3)]
+#     for i in range(args.n_layer):
+#         dd = model.strategy[i]
+#         dev = dd.device
+#         atype = dd.atype    
+#         state_init[i*3+0] = torch.zeros(args.n_embd, dtype=atype, requires_grad=False, device=dev).contiguous()
+#         state_init[i*3+1] = state_raw[f'blocks.{i}.att.time_state'].transpose(1,2).to(dtype=torch.float, device=dev).requires_grad_(False).contiguous()
+#         state_init[i*3+2] = torch.zeros(args.n_embd, dtype=atype, requires_grad=False, device=dev).contiguous()
+#     model_state = copy.deepcopy(state_init)
+
+def run_rnn(ctx, state):
+    ctx = ctx.replace("\r\n", "\n")
+    tokens = tokenizer.encode(ctx)
+    tokens = [int(x) for x in tokens]
+    
+    current_state = copy.deepcopy(state) if state is not None else None
+    
+    while len(tokens) > 0:
+        out, current_state = model.forward(tokens[:CHUNK_LEN], current_state)
+        tokens = tokens[CHUNK_LEN:]
+    
+    return out, current_state
+
+def generate_response(message, history, temperature=1.0, top_p=0.3):
+    global model_tokens, model_state
+    model_state = None
+    
+    ctx = ""
+    for human, assistant in history:
+        ctx += f"<|im_start|>user\n{human}<|im_end|>\n<|im_start|>assistant\n{assistant}<!--eos--><|im_end|>\n"
+    
+    ctx += f"<|im_start|>user\n{message}<|im_end|>\n<|im_start|>assistant\n"
+    
+    out, model_state = run_rnn(ctx, model_state)
+    
+    occurrence = {}
+    out_tokens = []
+    out_last = 0
+    response = ""
+    
+    eos_token_id = tokenizer.eos_token_id
+    im_end_id = tokenizer.encode("<|im_end|>")[0]  
+    for i in range(99999):
+        logits = out.clone()
+        for n in occurrence:
+            logits[n] -= GEN_alpha_presence + occurrence[n] * GEN_alpha_frequency
+        
+        logits[0] -= 1e10  
+        
+        token = sample_logits(logits, temperature=temperature, top_p=top_p)
+        
+        if token == im_end_id:
+            break
+            
+        out, model_state = model.forward([token], model_state)
+        
+        out_tokens += [token]
+        for xxx in occurrence:
+            occurrence[xxx] *= GEN_penalty_decay
+        occurrence[token] = 1 + (occurrence[token] if token in occurrence else 0)
+        
+        tmp = tokenizer.decode(out_tokens[out_last:])
+        if "\ufffd" not in tmp:  
+            response += tmp
+            cleaned_response = response.replace("<|im_end|>", "")
+            yield cleaned_response
+            out_last = i + 1
+            
+            if token == eos_token_id:
+                break
+
+def chat_with_bot(message, history, temperature, top_p):
+    response = ""
+    for partial_response in generate_response(message, history, temperature, top_p):
+        response = partial_response
+        yield response
+
+with gr.Blocks(title="MiniRWKV_7 DE 34.2M 🪿 2vGPU Space") as demo:
+    gr.Markdown("# MiniRWKV_7 DE 34.2M 🪿 ")
+    gr.Markdown("### Only 34.2M Params!!! Use 2V CPU Backend to run this model. ")
+    
+    with gr.Row():
+        with gr.Column(scale=3):
+            chatbot = gr.Chatbot(
+                label="对话记录",
+                height=1000,
+            )
+        
+        with gr.Column(scale=1):
+            msg = gr.Textbox(
+                label="输入消息",
+                placeholder="请输入您的问题...",
+                lines=3
+            )
+            
+            with gr.Row():
+                send_btn = gr.Button("发送", variant="primary")
+                clear_btn = gr.Button("清除历史")
+            
+            gr.Markdown("### 参数调节")
+            temperature_slider = gr.Slider(
+                minimum=0.1,
+                maximum=2.0,
+                value=GEN_TEMP,
+                step=0.1,
+                label="Temperature"
+            )
+            top_p_slider = gr.Slider(
+                minimum=0.0,
+                maximum=2.0,
+                value=GEN_TOP_P,
+                step=0.05,
+                label="Top-P"
+            )
+    
+    
+    def respond(message, chat_history, temperature, top_p):
+        if not message:
+            return "", chat_history
+        
+        chat_history.append((message, ""))
+        
+        response = ""
+        for partial_response in chat_with_bot(message, chat_history[:-1], temperature, top_p):
+            response = partial_response
+            cleaned_response = response.replace("<|im_end|>", "")
+            chat_history[-1] = (message, cleaned_response)
+            yield "", chat_history
+    
+    def clear_history():
+        global model_tokens, model_state
+        model_tokens = []
+        model_state = None
+        return []
+    
+    msg.submit(respond, [msg, chatbot, temperature_slider, top_p_slider], [msg, chatbot])
+    send_btn.click(respond, [msg, chatbot, temperature_slider, top_p_slider], [msg, chatbot])
+    clear_btn.click(clear_history, None, chatbot)
+
+if __name__ == "__main__":
+    demo.launch()

requirements.txt

@@ -0,0 +1,2 @@
+torch --index-url https://download.pytorch.org/whl/cpu
+transformers

sft-2048.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5c1e27de7156cb3c965a88409932805950e8839c998251183c7380390dd302fb
+size 182949245