上传 chat_server.py

2026-06-05 15:08:40 +00:00
1 parent 38b817a989
commit 390c980572
@@ -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, [])
+        
+        # 构建带上下文的输入（使用与训练一致的格式）
+        # 训练格式: <SOS>问题<SEP><START>
+        # 多轮格式: <SOS>历史问题1<SEP>历史回答1<SEP>历史问题2<SEP>历史回答2<SEP>当前问题<SEP><START>
+        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}<SEP>{a}")
+            
+            if context_parts:
+                context = "<SEP>".join(context_parts)
+                full_input = f"{context}<SEP>{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)