diff --git a/chat_server.py b/chat_server.py new file mode 100644 index 0000000..aa3a25e --- /dev/null +++ b/chat_server.py @@ -0,0 +1,143 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +""" +小狸聊天 Web 服务 v2 - 支持多轮对话 +""" +import sys, os, json, re +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) + +from flask import Flask, render_template, request, jsonify +import torch +from llm_chat_improved import ( + CharTokenizer, build_vocabulary, LLMModel, LLMChat, load_training_data +) + +app = Flask(__name__) + +# 全局模型 +llm_chat = None +vocab = None + +def clean_text(text): + """清洗文本""" + text = text.replace('\n', '') + text = re.sub(r'\*\*([^*]+)\*\*', r'\1', text) + text = re.sub(r'\([^)]*[ω。ノ♡^][^)]*\)', '', text) + emoji_pattern = re.compile( + '[\U0001F600-\U0001F64F\U0001F300-\U0001F5FF\U0001F680-\U0001F6FF\U0001F1E0-\U0001F1FF\U00002702-\U000027B0\U0001f926-\U0001f937\u2640-\u2642\u2600-\u2B55\u200d\u23cf\u23e9\u231a\ufe0f\u3030\u2728\u2714\u2716\u27a1\u2b05\u2b06\u2b07\u2b1b\u2b1c\u2b50\u2b55\u2934\u2935\u25aa\u25ab\u25b6\u25c0\u25fb\u25fc\u25fd\u25fe]+', + flags=re.UNICODE + ) + text = emoji_pattern.sub('', text) + text = re.sub(r'[ \t]+', ' ', text) + return text.strip() + +def load_model(): + """加载训练好的模型""" + global llm_chat, vocab + + print("🐱 正在加载小狸模型...") + + device = 'cpu' + + # 直接从模型文件加载 vocab 和配置 (v4大模型) + checkpoint = torch.load('xiaoli_final_v4.pth', map_location=device) + vocab = checkpoint['vocab'] + config = checkpoint.get('model_config', {}) + + print(f" 词汇表大小: {len(vocab)}") + + # 创建模型 - 使用保存的配置(v4大模型带MoE+Thinking+Memory) + model = LLMModel( + vocab_size=config.get('vocab_size', len(vocab)), + d_model=config.get('d_model', 384), + n_heads=config.get('n_heads', 8), + n_layers=config.get('n_layers', 6), + d_ff=config.get('d_ff', 1536), + max_len=config.get('max_len', 128), + dropout=0.0, + use_memory=config.get('use_memory', True), + use_thinking=config.get('use_thinking', True), + use_moe=config.get('use_moe', True), + num_experts=config.get('num_experts', 4), + top_k=config.get('top_k', 2), + ) + + # 加载权重 + model.load_state_dict(checkpoint['model_state_dict']) + + llm_chat = LLMChat(model, vocab, device, 128) + print("✅ 小狸已就绪!") + +# 存储会话历史 +session_history = {} + +@app.route('/') +def index(): + return render_template('chat.html') + +@app.route('/chat', methods=['POST']) +def chat(): + data = request.json + message = data.get('message', '').strip() + session_id = data.get('session_id', 'default') + + if not message: + return jsonify({'response': '请输入内容哦~'}) + + if llm_chat is None: + return jsonify({'response': '模型加载中,请稍后再试'}) + + try: + # 获取会话历史 + history = session_history.get(session_id, []) + + # 构建带上下文的输入(使用与训练一致的格式) + # 训练格式: 问题 + # 多轮格式: 历史问题1历史回答1历史问题2历史回答2当前问题 + if history and len(history) >= 2: + # 只保留最近2轮对话(4条记录) + recent_history = history[-4:] + # 提取问题和答案 + context_parts = [] + for i in range(0, len(recent_history), 2): + if i + 1 < len(recent_history): + q = recent_history[i].replace("用户:", "") + a = recent_history[i+1].replace("小狸:", "") + context_parts.append(f"{q}{a}") + + if context_parts: + context = "".join(context_parts) + full_input = f"{context}{message}" + else: + full_input = message + else: + full_input = message + + # 限制输入长度(字符数) + if len(full_input) > 80: + full_input = full_input[-80:] + + response = llm_chat.generate_response(full_input, max_length=50, temperature=0.3) + + # 更新历史 + history.append(f"用户:{message}") + history.append(f"小狸:{response}") + # 只保留最近6轮(12条记录) + session_history[session_id] = history[-12:] + + return jsonify({'response': response}) + except Exception as e: + import traceback + return jsonify({'response': f'出错了: {str(e)}\n{traceback.format_exc()}'}) + +@app.route('/clear', methods=['POST']) +def clear_history(): + data = request.json + session_id = data.get('session_id', 'default') + session_history[session_id] = [] + return jsonify({'status': 'ok'}) + +if __name__ == '__main__': + load_model() + print("\n🌐 启动 Web 服务: http://0.0.0.0:5000") + app.run(host='0.0.0.0', port=5000, debug=False)