Upload 4 files

优化推理时缓存管理与重复惩罚，添加历史记录输出！

app.py

@@ -1,294 +1,299 @@
-# 公开库
-import time
-import html
-import uuid
-import torch
-import threading
-import numpy as np
-import gradio as gr
-from queue import Queue
-# 私有库
-from make_model import make_model
-from LazyCache import ExpiringDict
-from train_and_use import El_text_continue_stream
-from tokenizer import tokenizer,vocab_size,token2str
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-# 加载模型
-model = make_model(
-    #token是从1开始的，0填充，剩下的用来覆盖全部字节
-    vocab_size = vocab_size+1+255,
-    embedding_dim = 768,
-    key_dim = 128,
-    head_number = 12,
-    position_information_type = "mask",
-    enable_affine = True,
-    enable_talking_head = True,
-    use_diff = False,
-    self_attention_block_size = 0,
-    feed_forward_dim = 1536,
-    enable_layer_norm = True,
-    deep = 12,
-    dropout_rate = 0.1,
-    enable_el_cache = True
-).to(device)
-model.load_state_dict(torch.load('large_model_instruct_09291732.weight',map_location=device,weights_only=True))
-model = model.eval()
-
-# token包装函数 - 使用HTML span标签确保每个token在独立矩形中
-def token_wapper(token):
-    # 对特殊字符进行HTML转义处理
-    escaped_token = html.escape(token)
-    return f'<span class="token-box">{escaped_token}</span>'
-
-# 多token包装函数 - 使用HTML span标签确保每个token在独立矩形中
-def token_split_wapper(token):
-    # 对特殊字符进行HTML转义处理
-    escaped_token = html.escape(token)
-    return f'<span class="multi-token-box">({escaped_token})[多token]</span>'
-
-# 处理用户输入的token，返回安全的显示格式
-def process_user_tokens(user_message):
-    # 通过分词器转化为token
-    user_tokens = tokenizer(user_message, 5.0)
-    
-    # 将token还原并进行安全包装
-    words = [] # token列表
-    temp = [] # token是特殊字节，要合并
-    for token in user_tokens:
-        if token > 0:
-            # 将合并成功的加入列表
-            if len(temp):
-                words.append(token_split_wapper(token2str(temp)))
-                temp = []
-            # 将新的token加入列表
-            words.append(token_wapper(token2str([token])))
-        else:
-            # 将字节送去合并
-            temp.append(token)
-    # 结束的时候要进行收尾
-    if len(temp):
-        words.append(token_split_wapper(token2str(temp)))
-    # 返回包装好的token列表
-    return ''.join(words)
-
-# 全局字典，存 per-session 的不可 deepcopy 对象 / 状态
-user_queues = ExpiringDict(ttl=550)             # session_id -> Queue(list([string,string]))，用于流式输出
-user_queues.start_auto_cleanup()
-user_stop_flags = ExpiringDict(ttl=550)         # session_id -> bool (True 表示停止)
-user_stop_flags.start_auto_cleanup()
-user_history_sessions_show = ExpiringDict(ttl=550)   # session_id -> 用于显示的历史记录，list([string,string])
-user_history_sessions_show.start_auto_cleanup()
-user_history_sessions_text = ExpiringDict(ttl=550)   # session_id -> 纯文本历史记录，string
-user_history_sessions_text.start_auto_cleanup()
-
-# 后台生成函数（只访问全局字典，通过 session_id 定位）
-def generate_text(sess, user_message, session_id, temperature, repeat_penalty, max_length, decay):
-    out = ""
-    q = user_queues.get(session_id)
-    # 立即刷出用户问题
-    q.put(out, block=False)
-    # 构建完整的对话历史输入
-    if len(sess) == 1:
-        user_history_sessions_text[session_id] = f"<|im_start|>user {user_message}<|im_end|><|im_start|>assistant "
-    else:
-        user_history_sessions_text[session_id] += f"<|im_start|>user {user_message}<|im_end|><|im_start|>assistant "
-    # 转换为模型输入格式
-    tokens_batch = [tokenizer(user_history_sessions_text[session_id], 5.0)]
-    tokens_batch = np.array(tokens_batch, dtype=np.int64) + 255
-    inputs = torch.from_numpy(tokens_batch).to(device).data
-    last_len = -1
-    # 模型输出
-    with torch.no_grad():
-        for o in El_text_continue_stream(
-            model, inputs, out_length=max_length,
-            repeat_penalty_value=repeat_penalty,
-            temperature=temperature,decay=decay,session_id=session_id):
-            # 如果当前位置可以完整解码
-            if o[0,-1] > 255:
-                # 将未解码的部分一起解码
-                temp = token2str(o[0][last_len:].cpu().numpy()-255)
-                out += temp
-                user_history_sessions_text[session_id] += temp
-                # 重置为解码光标
-                last_len = -1    
-                q.put(out, block=False)
-            else:
-                # 无法解码，光标固定
-                last_len -= 1
-            # 如果用户主动断开连接，停止生成,去除潜在标记
-            if user_stop_flags.get(session_id, True):
-                if '<' + out.split('<')[-1] in '<|im_end|>':
-                    # 显示的部分去除标记
-                    out = '<'+'<'.join(out.split('<')[:-1])
-                    # 历史的部分保留标记
-                    user_history_sessions_text[session_id] = '<'+'<'.join(user_history_sessions_text[session_id].split('<')[:-1])+'<|im_end|>'
-                break
-            
-            # 如果是输出终止标记
-            if '<|im_end|>' in out:
-                # 显示的部分，去除标记
-                out = out.split('<|im_end|>')[0]
-                q.put(out, block=False)
-                break
-            # 如果用户中断
-            if user_stop_flags[session_id] == True:
-                break
-        # 更新标记为暂停
-        user_stop_flags[session_id] = True
-
-# 按钮处理逻辑：发送消息 / 停止生成 / 清空会话
-def send_message(sess, btn_label, user_message, session_id, temperature, repeat_penalty, max_length, decay):
-    # 发送消息按钮 - 启动生成线程
-    if btn_label == "发送消息" and user_message:
-        # 设置当前用户正在生成的标志
-        user_stop_flags[session_id] = False
-        # 立即在UI中显示用户消息
-        user_tokens_display = process_user_tokens(user_message)
-        # 添加用户消息到当前会话
-        user_history_sessions_show[session_id] = sess
-        user_history_sessions_show[session_id] += [[user_tokens_display, ""]]
-        if session_id not in user_history_sessions_text:
-            return "", "会话过期！"
-        # 在这里开始流式输出
-        thread = threading.Thread(target=generate_text, args=(sess, user_message, session_id, temperature, repeat_penalty, max_length, decay))
-        thread.daemon = True #主进程退出时退出
-        thread.start() #启动
-        user_stop_flags[session_id] = False
-        # 清空输入框，更新按钮文本
-        return "", "停止生成"
-    else:
-        # 停止生成按钮 - 设置标志位
-        user_stop_flags[session_id] = True
-        # 更新返回给前端的 state/stop_flag
-        return user_message, "发送消息"
-
-# 清空会话
-def clear_session():
-    return []
-
-# 流式输出,无限循环刷新页面
-def stream_output(sess):
-    global user_queues, user_stop_flags, user_history_sessions_show, user_history_sessions_text
-    # 页面加载时初始化 session
-    session_id = str(uuid.uuid4())
-    user_queues[session_id] = Queue()
-    user_stop_flags[session_id] = True
-    user_history_sessions_show[session_id] = []  # 初始化历史会话记录，用于显示
-    user_history_sessions_text[session_id] = ""  # 初始化历史会话记录，用于文本存储
-    # 返回初始状态
-    yield [], "发送消息", session_id
-    # 不断刷新
-    while True:
-        time.sleep(0.01)  # 防止 busy-wait 占满 CPU
-        # 等待队列有数据
-        q = user_queues.get(session_id)
-        if q is None:
-            continue
-        # 处理队列中的消息
-        if not q.empty():
-            # 取到最后一个加入的数据
-            while q.qsize() > 1:
-                q.get()
-            out = q.get()
-            sess = user_history_sessions_show[session_id]
-            sess[-1][1] = out
-            # 更新UI状态
-            current_stopped = user_stop_flags.get(session_id, True)
-            button_label = "停止生成" if not current_stopped else "发送消息"
-            yield sess, button_label, session_id
-
-# UI美化
-css = """
-/* 大标题居中 */
-.title {
-    text-align: center;
-}
-/* 高级选项字体居中 */
-#adv-param button {
-    justify-content: center;
-}
-/* 高级选项字体放大 */
-#adv-param > button > span {
-    font-size: 16px !important;
-    font-weight: 600 !important;
-}
-/* 自定义token样式 */
-.token-box {
-    display: inline-block;
-    background-color: #f0f0f0;
-    border: 1px solid #ddd;
-    border-radius: 4px;
-    padding: 2px 4px;
-    margin: 2px;
-    font-family: monospace;
-}
-.multi-token-box {
-    display: inline-block;
-    background-color: #e6f7ff;
-    border: 1px solid #91d5ff;
-    border-radius: 4px;
-    padding: 2px 4px;
-    margin: 2px;
-    font-family: monospace;
-}
-"""
-# ========== Gradio UI ==========
-with gr.Blocks(css=css) as demo:
-    with gr.Column(elem_classes="container"):
-        gr.Markdown("# 0.18B中文大语言模型在线体验", elem_classes="title")
-        # 聊天界面
-        chatbot = gr.Chatbot(
-            label="对话", 
-            autoscroll=False,
-            show_copy_button=True,
-            elem_classes="chatbox",
-            type="tuples",
-            height=400
-        )
-        # 输入区域
-        with gr.Column(elem_classes="input-area"):
-            msg = gr.Textbox(
-                placeholder="请输入你的问题...", 
-                label="", 
-                lines=3,
-                show_label=False
-            )
-            # 按钮区域
-            with gr.Row(elem_classes="button-row"):
-                send_btn = gr.Button("发送消息", elem_classes="send-btn")
-                clear_btn = gr.Button("清空会话", elem_classes="clear-btn")
-        # 参数设置区域（可折叠）
-        with gr.Accordion("高级参数设置", open=False, elem_classes="parameter-row", elem_id="adv-param"):
-            with gr.Row():
-                temperature = gr.Slider(0.0001, 3.0001, value=0.0001, step=0.1, label="Temperature")
-                repeat_penalty = gr.Slider(0.0, 5.0, value=0.5, step=0.1, label="Repeat Penalty")
-            with gr.Row():
-                max_length = gr.Slider(64, 8192, value=512, step=64, label="Max Length")
-                decay = gr.Slider(0.90, 1.0, value=0.99, step=0.005, label="Repeat Penalty Decay Rate")
-    # gr.State 用来在前端保存可 deepcopied 的 session 值
-    session_id = gr.State()
-    # 发送按钮处理
-    send_btn.click(
-        send_message,
-        inputs=[chatbot, send_btn, msg, session_id, temperature, repeat_penalty, max_length, decay],
-        outputs=[msg, send_btn],
-    )
-    # 会话清空按钮处理
-    clear_btn.click(
-        clear_session,
-        inputs=[],
-        outputs=[chatbot],
-    )
-    # 无限循环，一直更新聊天界面
-    demo.load(
-        stream_output,
-        inputs=[chatbot],
-        outputs=[chatbot, send_btn, session_id],
-    )
-if __name__ == "__main__":
-    """主函数：启动Gradio界面"""
-    # 设置队列参数以提高并发处理能力
-    demo.queue(max_size=128, default_concurrency_limit=128)
-    # 启动Gradio应用，不公开分享，并应用CSS样式
-    demo.launch(share=False)
+# 公开库
+import time
+import html
+import uuid
+import torch
+import threading
+import numpy as np
+import gradio as gr
+from queue import Queue
+# 私有库
+from make_model import make_model
+from LazyCache import ExpiringDict
+from train_and_use import El_text_continue_stream
+from tokenizer import tokenizer,vocab_size,token2str
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# 加载模型
+model = make_model(
+    #token是从1开始的，0填充，剩下的用来覆盖全部字节
+    vocab_size = vocab_size+1+255,
+    embedding_dim = 768,
+    key_dim = 128,
+    head_number = 12,
+    position_information_type = "mask",
+    enable_affine = True,
+    enable_talking_head = True,
+    use_diff = False,
+    self_attention_block_size = 0,
+    feed_forward_dim = 1536,
+    enable_layer_norm = True,
+    deep = 12,
+    dropout_rate = 0.1,
+    enable_el_cache = True
+).to(device)
+model.load_state_dict(torch.load('large_model_instruct_09291732.weight',map_location=device,weights_only=True))
+model = model.eval()
+
+# token包装函数 - 使用HTML span标签确保每个token在独立矩形中
+def token_wapper(token):
+    # 对特殊字符进行HTML转义处理
+    escaped_token = html.escape(token)
+    return f'<span class="token-box">{escaped_token}</span>'
+
+# 多token包装函数 - 使用HTML span标签确保每个token在独立矩形中
+def token_split_wapper(token):
+    # 对特殊字符进行HTML转义处理
+    escaped_token = html.escape(token)
+    return f'<span class="multi-token-box">({escaped_token})[多token]</span>'
+
+# 处理用户输入的token，返回安全的显示格式
+def process_user_tokens(user_message):
+    # 通过分词器转化为token
+    user_tokens = tokenizer(user_message, 5.0)
+    
+    # 将token还原并进行安全包装
+    words = [] # token列表
+    temp = [] # token是特殊字节，要合并
+    for token in user_tokens:
+        if token > 0:
+            # 将合并成功的加入列表
+            if len(temp):
+                words.append(token_split_wapper(token2str(temp)))
+                temp = []
+            # 将新的token加入列表
+            words.append(token_wapper(token2str([token])))
+        else:
+            # 将字节送去合并
+            temp.append(token)
+    # 结束的时候要进行收尾
+    if len(temp):
+        words.append(token_split_wapper(token2str(temp)))
+    # 返回包装好的token列表
+    return ''.join(words)
+
+# 全局字典，存 per-session 的不可 deepcopy 对象 / 状态
+user_queues = ExpiringDict(ttl=550)             # session_id -> Queue(list([string,string]))，用于流式输出
+user_queues.start_auto_cleanup()
+user_stop_flags = ExpiringDict(ttl=550)         # session_id -> bool (True 表示停止)
+user_stop_flags.start_auto_cleanup()
+user_history_sessions_show = ExpiringDict(ttl=550)   # session_id -> 用于显示的历史记录，list([string,string])
+user_history_sessions_show.start_auto_cleanup()
+user_history_sessions_text = ExpiringDict(ttl=550)   # session_id -> 纯文本历史记录，string，有KV_Cache，用不到历史，但留着可作为日志输出[狗头]
+user_history_sessions_text.start_auto_cleanup()
+
+# 后台生成函数（只访问全局字典，通过 session_id 定位）
+def generate_text(sess, user_message, session_id, temperature, repeat_penalty, max_length, decay):
+    out = ""
+    q = user_queues.get(session_id)
+    # 立即刷出用户问题
+    q.put(out, block=False)
+    # 记录历史长度
+    history_len = len(sess) - 1
+    # 构建完整的对话历史输入
+    if history_len:
+        user_history_sessions_text[session_id] += f"<|im_start|>user {user_message}<|im_end|><|im_start|>assistant "
+    else:
+        user_history_sessions_text[session_id] = f"<|im_start|>user {user_message}<|im_end|><|im_start|>assistant "
+    # 转换为模型输入格式
+    tokens_batch = [tokenizer(f"<|im_start|>user {user_message}<|im_end|><|im_start|>assistant ", 5.0)]
+    tokens_batch = np.array(tokens_batch, dtype=np.int64) + 255
+    inputs = torch.from_numpy(tokens_batch).to(device).data
+    last_len = -1
+    # 模型输出
+    with torch.no_grad():
+        for o in El_text_continue_stream(
+            model, inputs, out_length=max_length,
+            repeat_penalty_value=repeat_penalty,
+            temperature=temperature,decay=decay,
+            session_id=session_id,history_len=history_len):
+            # 如果当前位置可以完整解码
+            if o[0,-1] > 255:
+                # 将未解码的部分一起解码
+                temp = token2str(o[0][last_len:].cpu().numpy()-255)
+                out += temp
+                user_history_sessions_text[session_id] += temp
+                # 重置为解码光标
+                last_len = -1    
+                q.put(out, block=False)
+            else:
+                # 无法解码，光标固定
+                last_len -= 1
+            # 如果用户主动断开连接，停止生成,去除潜在标记
+            if user_stop_flags.get(session_id, True):
+                if '<' + out.split('<')[-1] in '<|im_end|>':
+                    # 显示的部分去除标记
+                    out = '<'+'<'.join(out.split('<')[:-1])
+                    # 历史的部分保留标记
+                    user_history_sessions_text[session_id] = '<'+'<'.join(user_history_sessions_text[session_id].split('<')[:-1])+'<|im_end|>'
+                break
+            
+            # 如果是输出终止标记
+            if '<|im_end|>' in out:
+                # 显示的部分，去除标记
+                out = out.split('<|im_end|>')[0]
+                q.put(out, block=False)
+                break
+            # 如果用户中断
+            if user_stop_flags[session_id] == True:
+                break
+        # 更新标记为暂停
+        user_stop_flags[session_id] = True
+        # 打印日志用于查看用户提问，用于后续优化
+        print(user_history_sessions_text[session_id])
+
+# 按钮处理逻辑：发送消息 / 停止生成 / 清空会话
+def send_message(sess, btn_label, user_message, session_id, temperature, repeat_penalty, max_length, decay):
+    # 发送消息按钮 - 启动生成线程
+    if btn_label == "发送消息" and user_message:
+        # 设置当前用户正在生成的标志
+        user_stop_flags[session_id] = False
+        # 立即在UI中显示用户消息
+        user_tokens_display = process_user_tokens(user_message)
+        # 添加用户消息到当前会话
+        user_history_sessions_show[session_id] = sess
+        user_history_sessions_show[session_id] += [[user_tokens_display, ""]]
+        if session_id not in user_history_sessions_text:
+            return "", "会话过期！"
+        # 在这里开始流式输出
+        thread = threading.Thread(target=generate_text, args=(sess, user_message, session_id, temperature, repeat_penalty, max_length, decay))
+        thread.daemon = True #主进程退出时退出
+        thread.start() #启动
+        user_stop_flags[session_id] = False
+        # 清空输入框，更新按钮文本
+        return "", "停止生成"
+    else:
+        # 停止生成按钮 - 设置标志位
+        user_stop_flags[session_id] = True
+        # 更新返回给前端的 state/stop_flag
+        return user_message, "发送消息"
+
+# 清空会话
+def clear_session():
+    return []
+
+# 流式输出,无限循环刷新页面
+def stream_output(sess):
+    global user_queues, user_stop_flags, user_history_sessions_show, user_history_sessions_text
+    # 页面加载时初始化 session
+    session_id = str(uuid.uuid4())
+    user_queues[session_id] = Queue()
+    user_stop_flags[session_id] = True
+    user_history_sessions_show[session_id] = []  # 初始化历史会话记录，用于显示
+    user_history_sessions_text[session_id] = ""  # 初始化历史会话记录，用于文本存储
+    # 返回初始状态
+    yield [], "发送消息", session_id
+    # 不断刷新
+    while True:
+        time.sleep(0.01)  # 防止 busy-wait 占满 CPU
+        # 等待队列有数据
+        q = user_queues.get(session_id)
+        if q is None:
+            continue
+        # 处理队列中的消息
+        if not q.empty():
+            # 取到最后一个加入的数据
+            while q.qsize() > 1:
+                q.get()
+            out = q.get()
+            sess = user_history_sessions_show[session_id]
+            sess[-1][1] = out
+            # 更新UI状态
+            current_stopped = user_stop_flags.get(session_id, True)
+            button_label = "停止生成" if not current_stopped else "发送消息"
+            yield sess, button_label, session_id
+
+# UI美化
+css = """
+/* 大标题居中 */
+.title {
+    text-align: center;
+}
+/* 高级选项字体居中 */
+#adv-param button {
+    justify-content: center;
+}
+/* 高级选项字体放大 */
+#adv-param > button > span {
+    font-size: 16px !important;
+    font-weight: 600 !important;
+}
+/* 自定义token样式 */
+.token-box {
+    display: inline-block;
+    background-color: #f0f0f0;
+    border: 1px solid #ddd;
+    border-radius: 4px;
+    padding: 2px 4px;
+    margin: 2px;
+    font-family: monospace;
+}
+.multi-token-box {
+    display: inline-block;
+    background-color: #e6f7ff;
+    border: 1px solid #91d5ff;
+    border-radius: 4px;
+    padding: 2px 4px;
+    margin: 2px;
+    font-family: monospace;
+}
+"""
+# ========== Gradio UI ==========
+with gr.Blocks(css=css) as demo:
+    with gr.Column(elem_classes="container"):
+        gr.Markdown("# 0.18B中文大语言模型", elem_classes="title")
+        # 聊天界面
+        chatbot = gr.Chatbot(
+            label="对话", 
+            autoscroll=False,
+            show_copy_button=True,
+            elem_classes="chatbox",
+            type="tuples",
+            height=400
+        )
+        # 输入区域
+        with gr.Column(elem_classes="input-area"):
+            msg = gr.Textbox(
+                placeholder="请输入你的问题...", 
+                label="", 
+                lines=3,
+                show_label=False
+            )
+            # 按钮区域
+            with gr.Row(elem_classes="button-row"):
+                send_btn = gr.Button("发送消息", elem_classes="send-btn")
+                clear_btn = gr.Button("清空会话", elem_classes="clear-btn")
+        # 参数设置区域（可折叠）
+        with gr.Accordion("高级参数设置", open=False, elem_classes="parameter-row", elem_id="adv-param"):
+            with gr.Row():
+                temperature = gr.Slider(0.0001, 3.0001, value=0.0001, step=0.1, label="Temperature")
+                repeat_penalty = gr.Slider(0.0, 5.0, value=0.5, step=0.1, label="Repeat Penalty")
+            with gr.Row():
+                max_length = gr.Slider(64, 8192, value=512, step=64, label="Max Length")
+                decay = gr.Slider(0.90, 1.0, value=0.99, step=0.005, label="Repeat Penalty Decay Rate")
+    # gr.State 用来在前端保存可 deepcopied 的 session 值
+    session_id = gr.State()
+    # 发送按钮处理
+    send_btn.click(
+        send_message,
+        inputs=[chatbot, send_btn, msg, session_id, temperature, repeat_penalty, max_length, decay],
+        outputs=[msg, send_btn],
+    )
+    # 会话清空按钮处理
+    clear_btn.click(
+        clear_session,
+        inputs=[],
+        outputs=[chatbot],
+    )
+    # 无限循环，一直更新聊天界面
+    demo.load(
+        stream_output,
+        inputs=[chatbot],
+        outputs=[chatbot, send_btn, session_id],
+    )
+if __name__ == "__main__":
+    """主函数：启动Gradio界面"""
+    # 设置队列参数以提高并发处理能力
+    demo.queue(max_size=128, default_concurrency_limit=128)
+    # 启动Gradio应用，不公开分享，并应用CSS样式
+    demo.launch(share=False)

install_ac.py

@@ -0,0 +1,11 @@
+import os
+import sys
+import subprocess
+os.environ['AHOCORASICK_BYTES'] = '1'
+# 使用当前解释器对应的pip进行安装
+subprocess.check_call([
+    sys.executable, "-m", "pip", "install",
+    "--force-reinstall",  # 强制重新安装
+    "--no-cache-dir",     # 不使用缓存，确保获取最新版本
+    "git+https://github.com/WojciechMula/pyahocorasick.git"
+])

tokenizer.py

@@ -1,104 +1,91 @@
-import numpy as np
-import subprocess
-import sys
-import os
-
-# 设置环境变量
-os.environ['AHOCORASICK_BYTES'] = '1'
-# 使用当前解释器对应的pip进行安装
-subprocess.check_call([
-    sys.executable, "-m", "pip", "install",
-    "--force-reinstall",  # 强制重新安装
-    "--no-cache-dir",     # 不使用缓存，确保获取最新版本
-    "git+https://github.com/WojciechMula/pyahocorasick.git"
-])
-
-#加载词表
-with open('vocab_b_65544.txt','r',encoding='utf-8') as f:
-# with open('vocab_tiny_random.txt','r',encoding='utf-8') as f:
-    words_count = dict()
-    for word in f:
-        if word[0] != '\t':
-            k,v = word.split('\t')
-            words_count[k] = int(v[:-1])
-        
-#补充缺失词，但尽量不要改变词频
-if '.' in words_count:
-    words_count[','] = words_count['.']
-    words_count['\r'] = 1
-    words_count['\n'] = 1
-    words_count['\t'] = 1
-    
-#计算每个片段长度的单词总数
-N = 7
-count_sum = [0 for _ in range(N)]
-for k,v in words_count.items():
-    count_sum[len(k)-1] += v
-
-#创建AC自动机
-import ahocorasick as ah
-aca= ah.Automaton()
-for k,v in words_count.items():
-    aca.add_word(k.encode(),(len(k.encode()),np.log(v/count_sum[len(k)-1])))
-aca.make_automaton()
-
-#单词与整数互转字典
-words = [k for k in words_count]
-words.sort()
-word2idx = {k.encode():i+1 for i,k in enumerate(words)}
-idx2word = {i:k for k,i in word2idx.items()}
-vocab_size = len(word2idx)
-
-#分词器函数
-def tokenizer(text,alpha=1.0):
-    encode_text = text.encode()
-    #路径，记录起始位置和分值
-    LOT = len(encode_text)
-    BOW = 0 #表示最佳词的起始位置
-    VOW = 1 #表示最佳路径的累积值
-    VOID = 5 #表示没有记录
-    routes = [(i,VOID) for i in range(LOT)] + [(-1,0.0)]
-    tokens = []  #保存分词结果
-    #遍历所有匹配成功的词
-    # low:len_of_word
-    # vow:value_of_word
-    for eow, (low,vow) in aca.iter(encode_text):
-        #匹配词起点序号 = 匹配词终点序号 -（匹配词长度-1）
-        bow = eow - low + 1
-        #得分是负数，但负的程度越小约好，
-        #数值为起始位置的得分 + 当前词的分数
-        #起始位置无记录就往前找
-        i = 0
-        while routes[bow - 1 - i][VOW] == VOID:
-            i += 1
-        v = routes[bow - 1 -i][VOW] + vow
-        # 超过5.0直接使用确定算法
-        if alpha >= 5.0:
-            #更短的路径或第一个到达，更新
-            if v > routes[eow][VOW] and i == 0 or routes[eow][VOW] == VOID:
-                routes[eow] = bow,v #记录起始位置以及累积值
-        else:
-            # 随机算法
-            if routes[eow][VOW] == VOID:
-                base = v
-                temp = 1.0
-                denominator = 1.0
-                routes[eow] = bow,v
-            else:
-                temp = np.exp(alpha * (v - base))
-                denominator += temp
-                if np.random.rand() < temp/denominator and i == 0:
-                    routes[eow] = bow,v #记录起始位置以及累积值
-    #从后往前查找分割点
-    eow = LOT - 1
-    while encode_text:
-        bow = routes[eow][BOW] #找到最佳词的起始位置
-        tokens.append(encode_text[bow:eow+1]) #记录该词语
-        encode_text,eow = encode_text[:bow],bow - 1 #继续分上一个词
-    #从后往前找，需要反序得到正序的分词结果
-    return [word2idx[w] if w in word2idx else -ord(w) for w in tokens[::-1]]
-
-def token2str(tokens,split=''):
-    return b''.join([(int(-token)).to_bytes(1,'big') if token < 0 else idx2word[token] + split.encode() for token in tokens]).decode(errors="ignore")
-    
-
+import numpy as np
+import install_ac
+#加载词表
+with open('vocab_b_65544.txt','r',encoding='utf-8') as f:
+# with open('vocab_tiny_random.txt','r',encoding='utf-8') as f:
+    words_count = dict()
+    for word in f:
+        if word[0] != '\t':
+            k,v = word.split('\t')
+            words_count[k] = int(v[:-1])
+        
+#补充缺失词，但尽量不要改变词频
+if '.' in words_count:
+    words_count[','] = words_count['.']
+    words_count['\r'] = 1
+    words_count['\n'] = 1
+    words_count['\t'] = 1
+    
+#计算每个片段长度的单词总数
+N = 7
+count_sum = [0 for _ in range(N)]
+for k,v in words_count.items():
+    count_sum[len(k)-1] += v
+
+#创建AC自动机
+import ahocorasick as ah
+aca= ah.Automaton()
+for k,v in words_count.items():
+    aca.add_word(k.encode(),(len(k.encode()),np.log(v/count_sum[len(k)-1])))
+aca.make_automaton()
+
+#单词与整数互转字典
+words = [k for k in words_count]
+words.sort()
+word2idx = {k.encode():i+1 for i,k in enumerate(words)}
+idx2word = {i:k for k,i in word2idx.items()}
+vocab_size = len(word2idx)
+
+#分词器函数
+def tokenizer(text,alpha=1.0):
+    encode_text = text.encode()
+    #路径，记录起始位置和分值
+    LOT = len(encode_text)
+    BOW = 0 #表示最佳词的起始���置
+    VOW = 1 #表示最佳路径的累积值
+    VOID = 5 #表示没有记录
+    routes = [(i,VOID) for i in range(LOT)] + [(-1,0.0)]
+    tokens = []  #保存分词结果
+    #遍历所有匹配成功的词
+    # low:len_of_word
+    # vow:value_of_word
+    for eow, (low,vow) in aca.iter(encode_text):
+        #匹配词起点序号 = 匹配词终点序号 -（匹配词长度-1）
+        bow = eow - low + 1
+        #得分是负数，但负的程度越小约好，
+        #数值为起始位置的得分 + 当前词的分数
+        #起始位置无记录就往前找
+        i = 0
+        while routes[bow - 1 - i][VOW] == VOID:
+            i += 1
+        v = routes[bow - 1 -i][VOW] + vow
+        # 超过5.0直接使用确定算法
+        if alpha >= 5.0:
+            #更短的路径或第一个到达，更新
+            if v > routes[eow][VOW] and i == 0 or routes[eow][VOW] == VOID:
+                routes[eow] = bow,v #记录起始位置以及累积值
+        else:
+            # 随机算法
+            if routes[eow][VOW] == VOID:
+                base = v
+                temp = 1.0
+                denominator = 1.0
+                routes[eow] = bow,v
+            else:
+                temp = np.exp(alpha * (v - base))
+                denominator += temp
+                if np.random.rand() < temp/denominator and i == 0:
+                    routes[eow] = bow,v #记录起始位置以及累积值
+    #从后往前查找分割点
+    eow = LOT - 1
+    while encode_text:
+        bow = routes[eow][BOW] #找到最佳词的起始位置
+        tokens.append(encode_text[bow:eow+1]) #记录该词语
+        encode_text,eow = encode_text[:bow],bow - 1 #继续分上一个词
+    #从后往前找，需要反序得到正序的分词结果
+    return [word2idx[w] if w in word2idx else -ord(w) for w in tokens[::-1]]
+
+def token2str(tokens,split=''):
+    return b''.join([(int(-token)).to_bytes(1,'big') if token < 0 else idx2word[token] + split.encode() for token in tokens]).decode(errors="ignore")
+    
+

train_and_use.py

@@ -340,25 +340,37 @@ def El_text_continue(model,inputs,out_length,repeat_penalty_value,temperature,de
                 repeat_penalty[i][next_token[i]] -= repeat_penalty_value
         return inputs
 
-def El_text_continue_stream(model,inputs,out_length,repeat_penalty_value,temperature,decay=0.98,session_id='0'):
+def El_text_continue_stream(model,inputs,out_length,repeat_penalty_value,temperature,decay=0.98,session_id='0',history_len=0):
     if model.model_type == "generator":
-        assert len(inputs[0]) > 1, "初始序列长度必须大于1，与增量续写进行区分"
-        query = model.embedding(inputs)
-        prob_dist = model.projector(model.encoder(query,inputs==inputs,session_id)[:,-1,:])
-        repeat_penalty = torch.zeros_like(prob_dist, device=inputs.device)
-        for index in range(inputs.size(1)):
-            for line in range(inputs.size(0)):
-                repeat_penalty[line][inputs[line][index]] -= repeat_penalty_value
-            repeat_penalty *= decay
-        prob_dist += repeat_penalty
-        next_token = torch.multinomial(F.softmax(prob_dist/temperature, dim = -1), num_samples = 1)
-        inputs = torch.cat([inputs,next_token.to(inputs.device)], dim=-1)[:,-4:]
+        if history_len == 0:
+            # 没有历史记录就处理整个输入，提高并行性
+            assert len(inputs[0]) > 1, "初始序列长度必须大于1，与增量续写进行区分"
+            query = model.embedding(inputs)
+            prob_dist = model.projector(model.encoder(query,inputs==inputs,session_id)[:,-1,:])
+            repeat_penalty = torch.zeros_like(prob_dist, device=inputs.device)
+            # 不计算用户输入的重复惩罚
+            # for index in range(inputs.size(1)):
+            #     for line in range(inputs.size(0)):
+            #         repeat_penalty[line][inputs[line][index]] -= repeat_penalty_value
+            #     repeat_penalty *= decay
+            # prob_dist += repeat_penalty
+            next_token = torch.multinomial(F.softmax(prob_dist/temperature, dim = -1), num_samples = 1)
+            inputs = torch.cat([inputs,next_token.to(inputs.device)], dim=-1)[:,-4:]
+        else:
+            # 否则将用户的问题逐个token处理，减少运算量
+            query = model.embedding(inputs)
+            for i in range(query.size(1)-1):
+                model.encoder(query[:,i:i+1],(inputs[:,[-1]]==inputs[:,[-1]]),session_id)[:,-1,:]
+            prob_dist = model.projector(model.encoder(query[:,[-1]],(inputs[:,[-1]]==inputs[:,[-1]]),session_id)[:,-1,:])
+            repeat_penalty = torch.zeros_like(prob_dist, device=inputs.device)
+            next_token = torch.multinomial(F.softmax(prob_dist/temperature, dim = -1), num_samples = 1)
+            inputs = torch.cat([inputs,next_token.to(inputs.device)], dim=-1)[:,-4:]
         yield inputs
         for i in range(next_token.size(0)):
             repeat_penalty[i][next_token[i]] -= repeat_penalty_value
         for _ in range(0,out_length-1,1):
             query = model.embedding(inputs[:,[-1]])
-            prob_dist = model.projector(model.encoder(query,(inputs==inputs)[:,[-1]],session_id)[:,-1,:])
+            prob_dist = model.projector(model.encoder(query,(inputs[:,[-1]]==inputs[:,[-1]]),session_id)[:,-1,:])
             repeat_penalty *= decay
             prob_dist += repeat_penalty
             next_token = torch.multinomial(F.softmax(prob_dist/temperature, dim = -1), num_samples = 1)