########################################################################################################
# The RWKV Language Model - https://github.com/BlinkDL/RWKV-LM
########################################################################################################
 
print("RWKV Chat Simple Demo")  # 打印一个简单的消息，表明这是 RWKV 聊天的简单演示。
import os, copy, types, gc, sys, re  # 导入操作系统、对象复制、类型、垃圾回收、系统、正则表达式等包
import numpy as np  # 导入 numpy 库
from prompt_toolkit import prompt  # 从 prompt_toolkit 导入 prompt，用于命令行输入
import torch  # 导入 pytorch 库

# 优化 PyTorch 设置，允许使用 tf32 
torch.backends.cudnn.benchmark = True
torch.backends.cudnn.allow_tf32 = True
torch.backends.cuda.matmul.allow_tf32 = True
 
os.environ["RWKV_JIT_ON"] = "1" # 启用 JIT 编译
os.environ["RWKV_CUDA_ON"] = "0"  # 禁用原生 CUDA 算子，改成 '1' 表示启用 CUDA 算子（速度更快，但需要 c++ 编译器和 CUDA 库）

from rwkv.model import RWKV  # 从 RWKV 模型库中导入 RWKV 类，用于加载和操作 RWKV 模型。
from rwkv.utils import PIPELINE  # 从 RWKV 工具库中导入 PIPELINE，用于数据的编码和解码
 
args = types.SimpleNamespace()
 
args.strategy = "cuda fp16"  # 模型推理的设备和精度，使用 CUDA （GPU）并采用 FP16 精度
args.MODEL_NAME = "E://RWKV-Runner//models//rwkv-final-v6-2.1-1b6"  # 指定 RWKV 模型的路径，建议写绝对路径

# STATE_NAME = None # 不使用 State
 
# 指定要加载的 State 文件路径。
STATE_NAME = "E://RWKV-Runner//models//rwkv-x060-eng_single_round_qa-1B6-20240516-ctx2048"  # 指定要加载的自定义 State 文件路径。

# 设置模型的解码参数
GEN_TEMP = 1.0
GEN_TOP_P = 0.3
GEN_alpha_presence = 0.5
GEN_alpha_frequency = 0.5
GEN_penalty_decay = 0.996
 
# 判断是否加载了一个 State 文件。如果指定了 State ，则调整生成参数，使回答的效果更好。
if STATE_NAME != None:
    GEN_TOP_P = 0.2
    GEN_alpha_presence = 0.3
    GEN_alpha_frequency = 0.3
 
CHUNK_LEN = 256  # 对输入进行分块处理

print(f"Loading model - {args.MODEL_NAME}")# 打印模型的加载消息
model = RWKV(model=args.MODEL_NAME, strategy=args.strategy)  # 加载 RWKV 模型。
pipeline = PIPELINE(model, "rwkv_vocab_v20230424")  # 初始化 PIPELINE ，使用 RWKV-World 词表处理输入和输出的编码/解码。

model_tokens = []
model_state = None
 
# 如果指定了 STATE_NAME，则加载自定义 State 文件，并初始化模型 State
if STATE_NAME != None:
    args = model.args  # 获取模型参数
    state_raw = torch.load(STATE_NAME + '.pth')  # 从指定的 State 文件中加载 State 数据
    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  # 获取每一层的加载设备（如 GPU）
        atype = dd.atype  # 获取每一层的数据类型（FP32/FP16 或 int8 等）
        # 初始化模型的状态
        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):
    # 定义两个全局变量，用于更新 token 和模型状态（state）
    global model_tokens, model_state
    ctx = ctx.replace("\r\n", "\n")  # 将文本中的 CRLF（Windows 系统的换行符）转换为 LF（Linux 系统的换行符）
    tokens = pipeline.encode(ctx)  # 基于 RWKV 模型的词汇表，将文本编码为 tokens
    tokens = [int(x) for x in tokens]  # 将 tokens 转换为整数（int）列表，确保类型一致性
    model_tokens += tokens  # 将 tokens 添加到全局的模型 token 列表中
 
    while len(tokens) > 0:  # 使用一个 while 循环执行模型前向传播，直到所有 tokens 处理完毕
        out, model_state = model.forward(tokens[:CHUNK_LEN], model_state)  # 模型前向传播，处理大小为 CHUNK_LEN 的 token 列表，并更新模型状态
        tokens = tokens[CHUNK_LEN:]  # 移除已处理的 tokens 块，并继续处理剩余的 tokens
 
    return out  # 返回模型的 prefill 结果

# 如果没有加载自定义 State ，则使用初始提示进行对话
if STATE_NAME == None:
    init_ctx = "User: hi" + "\n\n"
    init_ctx += "Assistant: Hi. I am your assistant and I will provide expert full response in full details. Please feel free to ask any question and I will always answer it." + "\n\n"
    run_rnn(init_ctx)  # 运行 RNN 模式对初始提示文本进行 prefill
    print(init_ctx, end="")  # 打印初始化对话文本

# 从用户输入中读取消息、循环生成下一个 token 
while True:
    msg = prompt("User: ")  # 从用户输入中读取消息，存到 msg 变量
    msg = msg.strip()  # 使用 strip 方法去除消息的首尾空格
    msg = re.sub(r"\n+", "\n", msg)  # 替换多个换行符为单个换行符
    if len(msg) > 0:  # 如果处理完后，用户输入的消息非空
        occurrence = {}  # 使用 occurrence 字典这个字典用于记录每个 token 在生成上下文中出现的次数，等会用在实现重复惩罚（Penalty）
        out_tokens = []  # 使用 out_tokens 列表记录即将输出的 tokens
        out_last = 0  # 用于记录上一次生成的 token 位置
 
        out = run_rnn("User: " + msg + "\n\nAssistant:")  # 将用户输入拼接成 RWKV 数据集的对话格式，进行 prefill  
        print("\nAssistant:", end="")  # 打印 "Assistant:" 标签
 
        for i in range(99999):  
            for n in occurrence: 
                out[n] -= GEN_alpha_presence + occurrence[n] * GEN_alpha_frequency  # 应用存在惩罚和频率惩罚参数
            out[0] -= 1e10  # 禁用 END_OF_TEXT 
 
            token = pipeline.sample_logits(out, temperature=GEN_TEMP, top_p=GEN_TOP_P)  # 采样生成下一个 token
 
            out, model_state = model.forward([token], model_state)  # 模型前向传播
            model_tokens += [token] 
            out_tokens += [token]  # 将新生成的 token 添加到输出的 token 列表中
 
            for xxx in occurrence:
                occurrence[xxx] *= GEN_penalty_decay  # 应用衰减重复惩罚
            occurrence[token] = 1 + (occurrence[token] if token in occurrence else 0)  # 更新 token 的出现次数
 
            tmp = pipeline.decode(out_tokens[out_last:])  # 将最新生成的 token 解码成文本
            if ("\ufffd" not in tmp) and (not tmp.endswith("\n")):  # 当生成的文本是有效 UTF-8 字符串且不以换行符结尾时
                print(tmp, end="", flush=True) #实时打印解码得到的文本
                out_last = i + 1 #更新输出位置变量 out_last 
 
            if "\n\n" in tmp:  # 如果生成的文本包含双换行符，表示模型的响应已结束（可以将 \n\n 改成其他停止词）
                print(tmp, end="", flush=True) # 实时打印解码得到的文本
                break #结束本轮推理
    else:
        print("!!! Error: please say something !!!")  # 如果用户没有输入消息，提示“输入错误，说点啥吧！”