144 lines
4.9 KiB
Python
144 lines
4.9 KiB
Python
#!/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)
|