Create app.py

app.py

@@ -0,0 +1,412 @@
+########################################################################################################
+# The RWKV Language Model - https://github.com/BlinkDL/RWKV-LM
+########################################################################################################
+
+import torch.nn.functional as F
+import gradio as gr
+from huggingface_hub import hf_hub_download
+import gc
+from rwkv.utils import PIPELINE, PIPELINE_ARGS
+import types, torch, copy, time
+from typing import List
+# torch.backends.cudnn.benchmark = True
+# torch.backends.cudnn.allow_tf32 = True
+# torch.backends.cuda.matmul.allow_tf32 = True
+# torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = True
+# torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = True
+torch._C._jit_set_autocast_mode(False)
+
+import torch.nn as nn
+from torch.nn import functional as F
+
+MyModule = torch.jit.ScriptModule
+MyFunction = torch.jit.script_method
+MyStatic = torch.jit.script
+pipeline = PIPELINE(None, "rwkv_vocab_v20230424")
+
+########################################################################################################
+
+print('\nNOTE: this is very inefficient (loads all weights to VRAM, and slow KV cache). better method is to prefetch DeepEmbed from RAM/SSD\n')
+
+args = types.SimpleNamespace()
+model_path = hf_hub_download(repo_id='Alic-Li/RWKV_v7_G1_Translate_ctx4096_20250620', filename='RWKV_v7s_G1_DEA_0.1B_Translate_ctx4096_20250917_latest.pth')
+args.MODEL_NAME = model_path
+args.n_layer = 12
+args.n_embd = 768
+args.vocab_size = 65536
+args.head_size = 64
+ctx_limit = 4096
+gen_limit = 4096
+penalty_decay = 0.996
+NUM_TRIALS = 1
+LENGTH_PER_TRIAL = 500
+TEMPERATURE = 1.0
+TOP_P = 0.0
+DTYPE = torch.half
+from torch.utils.cpp_extension import load
+HEAD_SIZE = args.head_size
+# ROCm_flag = torch.version.hip is not None 
+# if ROCm_flag:
+#     load(name="wkv7s", sources=["cuda/wkv7s_op.cpp", f"cuda/wkv7s.cu"], is_python_module=False,
+#                     verbose=True, extra_cuda_cflags=["-xhip", "-fopenmp", "-ffast-math", "-O3", "-munsafe-fp-atomics", f"-D_N_={HEAD_SIZE}"])
+# else:
+#     load(name="wkv7s", sources=["cuda/wkv7s_op.cpp", f"cuda/wkv7s.cu"], is_python_module=False,
+#                     verbose=True, extra_cuda_cflags=["-res-usage", "--use_fast_math", "-O3", "-Xptxas -O3", "--extra-device-vectorization", f"-D_N_={HEAD_SIZE}"])
+class WKV_7(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, state, r, w, k, v, a, b):
+        with torch.no_grad():
+            T, C = r.size()
+            H = C // HEAD_SIZE
+            N = HEAD_SIZE
+            assert HEAD_SIZE == C // H
+            assert all(x.dtype == DTYPE for x in [r,w,k,v,a,b])
+            assert all(x.is_contiguous() for x in [r,w,k,v,a,b])
+            y = torch.empty((T, C), device=k.device, dtype=DTYPE, requires_grad=False, memory_format=torch.contiguous_format)
+            torch.ops.wkv7s.forward(1, T, C, H, state, r, w, k, v, a, b, y)
+            return y
+def RWKV7_OP(state, r, w, k, v, a, b):
+    return WKV_7.apply(state, r, w, k, v, a, b)
+
+########################################################################################################
+
+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='cpu')
+        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 or 'qq.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[f'blocks.{i}.qkv.k_emb.weight'] = z[f'blocks.{i}.qkv.k_emb.weight'] + z['emb.weight'] @ z[f'blocks.{i}.qkv.k_emb_x.weight'].t()
+            z[f'blocks.{i}.qkv.v_emb.weight'] = z[f'blocks.{i}.qkv.v_emb.weight'] + z['emb.weight'] @ z[f'blocks.{i}.qkv.v_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 + 37)] # with KV cache etc.
+            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="cpu")
+                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="cpu")
+                state[i*3+2] = torch.zeros(args.n_embd, dtype=DTYPE, requires_grad=False, device="cpu")
+            state[args.n_layer*3+0] = torch.empty((0), dtype=torch.int, requires_grad=False, device="cpu") # token idx cache
+            for i in range(1,1+24): # kv cache = 12*2*32 numbers per token
+                state[args.n_layer*3+i] = torch.empty((0,32), dtype=DTYPE, requires_grad=False, device="cpu")
+            for i in range(1+24,1+36): # token-shift cache for Q in DEA
+                state[args.n_layer*3+i] = torch.zeros(256, dtype=DTYPE, requires_grad=False, device="cpu")
+
+        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) # sorry too busy to add forward_one mode
+                return self.forward_seq(idx, state, full_output)
+        else:
+            # return self.forward_one(idx, state) # sorry too busy to add forward_one mode
+            return self.forward_seq([idx], state, full_output)
+        
+    @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]
+            state[self.n_layer*3] = torch.cat((state[self.n_layer*3], torch.tensor(idx, dtype=torch.int, device=x.device)), dim=0)
+            ctx = state[self.n_layer*3]
+
+            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.'
+
+                qkv = f'blocks.{i}.qkv.'
+                q = x @ z[qkv+'qq.weight']                
+                k = x @ z[qkv+'k1']
+                state[self.n_layer*3+1+i*2] = torch.cat((state[self.n_layer*3+1+i*2], k), dim=0)
+                k = (state[self.n_layer*3+1+i*2] @ z[qkv+'k2']) * (z[qkv+'k_emb.weight'][ctx])
+                v = x @ z[qkv+'v1']
+                state[self.n_layer*3+1+i*2+1] = torch.cat((state[self.n_layer*3+1+i*2+1], v), dim=0)                        
+                v = torch.tanh(state[self.n_layer*3+1+i*2+1] @ z[qkv+'v2']) * (z[qkv+'v_emb.weight'][ctx])
+                qq = torch.cat((state[self.n_layer*3+1+24+i].unsqueeze(0), q[:-1,:]))
+                state[self.n_layer*3+1+24+i] = q[-1,:]
+                q = q + (qq - q) * z[qkv+'x_q']
+                k = k + (F.pad(k,(0, 0, 1, -1)) - k) * z[qkv+'x_k']
+                v = v + (F.pad(v,(0, 0, 1, -1)) - v) * z[qkv+'x_v']
+                q = F.layer_norm(q, (256,), weight=z[qkv+'lnq.weight'], bias=z[qkv+'lnq.bias'])
+                k = F.layer_norm(k, (256,), weight=z[qkv+'lnk.weight'], bias=z[qkv+'lnk.bias'])
+                v = F.layer_norm(v, (self.n_embd,), weight=z[qkv+'lnv.weight'], bias=z[qkv+'lnv.bias'])
+                scores = 64 * torch.tanh((q @ k.mT) * (1.0 / 1024.0)) # using soft-cap
+                if len(idx) > 1:
+                    mask = ~torch.tril(torch.ones(len(ctx), len(ctx), dtype=torch.bool, device=x.device))[-len(idx):,:]
+                    scores = scores.masked_fill(mask, float('-inf'))
+                qkv = scores.softmax(dim=-1) @ v
+
+                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 + qkv
+
+                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_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_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,:]
+
+model = RWKV_x070(args)
+
+def evaluate(
+    ctx,
+    token_count=200,
+    temperature=0.0,
+    top_p=0.0,
+    presencePenalty = 0.0,
+    countPenalty = 0.0,
+):
+    args = PIPELINE_ARGS(temperature = max(0.2, float(temperature)), top_p = float(top_p),
+                     alpha_frequency = countPenalty,
+                     alpha_presence = presencePenalty,
+                     token_ban = [], # ban the generation of some tokens
+                     token_stop = [0]) # stop generation whenever you see any token here
+    ctx = ctx.strip()
+    all_tokens = []
+    out_last = 0
+    out_str = ''
+    occurrence = {}
+    state = None
+    for i in range(int(token_count)):
+
+        input_ids = pipeline.encode(ctx)[-ctx_limit:] if i == 0 else [token]
+        out, state = model.forward(input_ids, state)
+        for n in occurrence:
+            out[n] -= (args.alpha_presence + occurrence[n] * args.alpha_frequency)
+
+        token = pipeline.sample_logits(out, temperature=args.temperature, top_p=args.top_p)
+        if token in args.token_stop:
+            break
+        all_tokens += [token]
+        for xxx in occurrence:
+            occurrence[xxx] *= penalty_decay
+            
+        ttt = pipeline.decode([token])
+        www = 1
+        if ttt in ' \t0123456789':
+            www = 0
+        #elif ttt in '\r\n,.;?!"\':+-*/=#@$%^&_`~|<>\\()[]{}，。；“”：？！（）【】':
+        #    www = 0.5
+        if token not in occurrence:
+            occurrence[token] = www
+        else:
+            occurrence[token] += www
+            
+        tmp = pipeline.decode(all_tokens[out_last:])
+        if '\ufffd' not in tmp:
+            out_str += tmp
+            yield out_str.strip()
+            out_last = i + 1
+    del out
+    del state
+    gc.collect()
+    # torch.cuda.empty_cache()
+    yield out_str.strip()
+
+def translate_english_to_chinese(english_text, token_count, temperature, top_p, presence_penalty, count_penalty):
+    if not english_text.strip():
+        return "Chinese:\n请输入英文内容。"
+    
+    full_prompt = f"English: {english_text}\n\nChinese:"
+    for output in evaluate(full_prompt, token_count, temperature, top_p, presence_penalty, count_penalty):
+        yield output  
+
+def translate_chinese_to_chinses(Chinese_text, token_count, temperature, top_p, presence_penalty, count_penalty):
+    if not Chinese_text.strip():
+        return "Chinses:\n请输入中文内容。"
+    
+    full_prompt = f"Chinese: {Chinese_text}\n\nEnglish:"
+    for output in evaluate(full_prompt, token_count, temperature, top_p, presence_penalty, count_penalty):
+        yield output  
+
+
+with gr.Blocks(title="RWKV_v7s_G1_DEA_0.1B_Translate_ctx4096_20250917 English -> Chinese") as demo:
+    with gr.Tab("English To Chinses"):    
+        gr.HTML(f"<div style='text-align:center;'><h1>RWKV_v7s_G1_DEA_0.1B_Translate_ctx4096_20250917 English -> Chinese</h1></div>")
+        with gr.Row():
+            with gr.Column():
+                english_input = gr.Textbox(
+                    label="英文输入(注意不能有空行)",
+                    lines=20,
+                    placeholder="请输入英文内容...",
+                    value="ROCm is an open-source stack, composed primarily of open-source software, designed for graphics processing unit (GPU) computation. ROCm consists of a collection of drivers, development tools, and APIs that enable GPU programming from low-level kernel to end-user applications.\n"
+                        "With ROCm, you can customize your GPU software to meet your specific needs.You can develop, collaborate, test, and deploy your applications in a free, open source, integrated, and secure software ecosystem. ROCm is particularly well-suited to GPU-accelerated high-performance computing (HPC), artificial intelligence (AI), scientific computing, and computer aided design (CAD).\n"
+                        "ROCm is powered by AMD’s Heterogeneous-computing Interface for Portability (HIP), an open-source software C++ GPU programming environment and its corresponding runtime. HIP allows ROCm developers to create portable applications on different platforms by deploying code on a range of platforms, from dedicated gaming GPUs to exascale HPC clusters.\n"
+                        "ROCm supports programming models, such as OpenMP and OpenCL, and includes all necessary open source software compilers, debuggers, and libraries. ROCm is fully integrated into machine learning (ML) frameworks, such as PyTorch and TensorFlow."
+                )
+
+            with gr.Column():
+                chinese_output = gr.Textbox(
+                    label="中文输出",
+                    lines=20,
+                    placeholder="翻译结果将显示在此处",
+                    value=""
+                )
+
+        with gr.Row():
+            translate_btn = gr.Button("Translate", variant="primary")
+            clear_btn = gr.Button("Clear", variant="secondary")
+            stop_btn = gr.Button("Stop", variant="stop")
+
+        with gr.Accordion("Advanced Settings", open=False):
+            token_count = gr.Slider(10, gen_limit, label="Max Tokens", step=10, value=gen_limit)
+            temperature = gr.Slider(0.2, 2.0, label="Temperature", step=0.1, value=1.0)
+            top_p = gr.Slider(0.0, 1.0, label="Top P", step=0.05, value=0)
+            presence_penalty = gr.Slider(0.0, 1.0, label="Presence Penalty", step=0.1, value=0)
+            count_penalty = gr.Slider(0.0, 1.0, label="Count Penalty", step=0.1, value=0)
+
+        translate_event = translate_btn.click(
+            fn=translate_english_to_chinese,
+            inputs=[english_input, token_count, temperature, top_p, presence_penalty, count_penalty],
+            outputs=[chinese_output]
+        )
+
+        clear_btn.click(
+            fn=lambda: ("", ""),
+            inputs=[],
+            outputs=[english_input, chinese_output]
+        )
+
+        stop_btn.click(
+            fn=None,
+            inputs=None,
+            outputs=None,
+            cancels=[translate_event]  
+        )
+    with gr.Tab("Chinses To English"):    
+        gr.HTML(f"<div style='text-align:center;'><h1>RWKV_v7s_G1_DEA_0.1B_Translate_ctx4096_20250917 Chinses -> English</h1></div>")
+        with gr.Row():
+            with gr.Column():
+                chinese_input = gr.Textbox(
+                    label="中文输入(注意不能有空行)",
+                    lines=20,
+                    placeholder="请输入中文内容...",
+                    value="ROCm是一个开源栈，主要由开源软件组成，旨在用于图形处理单元（GPU）计算。ROCm由一系列驱动程序、开发工具和API组成，这些工具和API允许从低级内核到最终用户应用程序对GPU进行编程。"
+                        "使用ROCm，您可以根据您的特定需求定制GPU软件。您可以在一个免费、开源、集成和安全的软件生态系统中开发、协作、测试和部署应用程序。ROCm特别适合GPU加速的高性能计算（HPC）、人工智能（AI）、科学计算和计算机辅助设计（CAD）。"
+                        "ROCm由AMD的可移植性图形处理接口（HIP）驱动，这是一个开源的C++ GPU编程环境及其相应的运行时。HIP允许ROCm开发者在不同平台上创建可移植应用程序，通过在从专用游戏GPU到exascale HPC集群的各种平台上部署代码来实现这一目标。"
+                        "ROCm支持编程模型，如OpenMP和OpenCL，并包含所有必要的开源软件编译器、调试器和库。ROCm完全集成到机器学习（ML）框架中，如PyTorch和TensorFlow。"
+                )
+
+            with gr.Column():
+                english_output = gr.Textbox(
+                    label="英文输出",
+                    lines=20,
+                    placeholder="翻译结果将显示在此处",
+                    value=""
+                )
+
+        with gr.Row():
+            translate_btn = gr.Button("Translate", variant="primary")
+            clear_btn = gr.Button("Clear", variant="secondary")
+            stop_btn = gr.Button("Stop", variant="stop")
+
+        with gr.Accordion("Advanced Settings", open=False):
+            token_count = gr.Slider(10, gen_limit, label="Max Tokens", step=10, value=gen_limit)
+            temperature = gr.Slider(0.2, 2.0, label="Temperature", step=0.1, value=1.0)
+            top_p = gr.Slider(0.0, 1.0, label="Top P", step=0.05, value=0)
+            presence_penalty = gr.Slider(0.0, 1.0, label="Presence Penalty", step=0.1, value=0)
+            count_penalty = gr.Slider(0.0, 1.0, label="Count Penalty", step=0.1, value=0)
+
+        translate_event = translate_btn.click(
+            fn=translate_chinese_to_chinses,
+            inputs=[chinese_input, token_count, temperature, top_p, presence_penalty, count_penalty],
+            outputs=[english_output]
+        )
+
+        clear_btn.click(
+            fn=lambda: ("", ""),
+            inputs=[],
+            outputs=[chinese_input, english_output]
+        )
+
+        stop_btn.click(
+            fn=None,
+            inputs=None,
+            outputs=None,
+            cancels=[translate_event]  
+        )
+
+demo.queue(max_size=10, default_concurrency_limit=1)
+demo.launch()