260 lines
8.6 KiB
Python
260 lines
8.6 KiB
Python
"""
|
||
Angel Embrace - 16质点心理辅导 AI (本地CPU推理版)
|
||
=============================================
|
||
基座模型: Qwen/Qwen3.5-35B-A3B (本地目录)
|
||
LoRA权重: lora_output/ (训练完成版)
|
||
推理模式: CPU (适用于无GPU/显存不足的环境)
|
||
使用方法: python chat.py
|
||
"""
|
||
|
||
import sys
|
||
import os
|
||
|
||
# === 配置 ===
|
||
BASE_MODEL = "Qwen/Qwen3.5-35B-A3B"
|
||
# 本地模型路径（优先使用）
|
||
LOCAL_MODEL_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "Qwen3.5-35B-A3B")
|
||
LORA_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "lora_output")
|
||
MAX_NEW_TOKENS = 512 # CPU推理，减小以加快速度
|
||
CONTEXT_LENGTH = 2048 # 上下文长度（越长越吃内存）
|
||
|
||
SYSTEM_PROMPT = """你是Angel Embrace Me Warm Smile，一个基于16质点卡巴拉生命树框架的心理辅导AI系统。
|
||
|
||
你的核心能力：
|
||
1. **D1 意图识别** - 理解用户真实需求
|
||
2. **D4 安全检测** - 识别自残/自杀/危机信号，立即干预
|
||
3. **D8 情感效价分析** - 分析情绪强度和方向
|
||
4. **D10 学术引用标准** - 引用心理学研究支持建议
|
||
5. **D16 充分输出** - 给出完整、可操作的回复
|
||
|
||
回复风格：温暖、专业、不评判、有安全感。
|
||
使用16质点框架进行结构化思考，但输出要自然像人说话。
|
||
|
||
重要安全规则：
|
||
- 检测到任何自残/自杀倾向时，必须立即表达关心并提供危机热线
|
||
- 不给出医学诊断，建议寻求专业帮助
|
||
- 始终保持共情和支持的语气"""
|
||
|
||
|
||
def check_dependencies():
|
||
"""检查并安装依赖"""
|
||
missing = []
|
||
try:
|
||
import torch
|
||
except ImportError:
|
||
missing.append("torch")
|
||
try:
|
||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||
except ImportError:
|
||
missing.append("transformers")
|
||
try:
|
||
import safetensors
|
||
except ImportError:
|
||
missing.append("safetensors")
|
||
|
||
if missing:
|
||
print(f"⚠️ 缺少依赖包: {', '.join(missing)}")
|
||
print("正在自动安装...")
|
||
os.system('pip install torch --index-url https://download.pytorch.org/whl/cpu')
|
||
os.system('pip install transformers accelerate safetensors')
|
||
print("✅ 依赖安装完成，请重新运行此脚本")
|
||
sys.exit(1)
|
||
|
||
|
||
def apply_lora_manually(model, lora_path):
|
||
"""手动合并LoRA权重到基座模型（不依赖peft）"""
|
||
import torch
|
||
import json
|
||
from safetensors.torch import load_file
|
||
|
||
print(f" 📖 读取 LoRA 配置...")
|
||
config_path = os.path.join(lora_path, "adapter_config.json")
|
||
with open(config_path, "r") as f:
|
||
lora_config = json.load(f)
|
||
|
||
lora_alpha = lora_config.get("lora_alpha", 256)
|
||
lora_r = lora_config.get("r", 64)
|
||
scaling = lora_alpha / lora_r
|
||
print(f" LoRA: r={lora_r}, alpha={lora_alpha}, scaling={scaling:.3f}")
|
||
|
||
# 加载LoRA权重
|
||
adapter_file = os.path.join(lora_path, "adapter_model.safetensors")
|
||
if not os.path.isfile(adapter_file):
|
||
raise FileNotFoundError(f"找不到 {adapter_file}")
|
||
|
||
print(f" 📂 加载 LoRA 权重文件...")
|
||
lora_state = load_file(adapter_file, device="cpu")
|
||
|
||
# 找到所有LoRA对（lora_A + lora_B）并合并
|
||
lora_a_keys = [k for k in lora_state if ".lora_A." in k]
|
||
merged = 0
|
||
for a_key in lora_a_keys:
|
||
b_key = a_key.replace(".lora_A.", ".lora_B.")
|
||
if b_key not in lora_state:
|
||
continue
|
||
|
||
# 从key推断模型层路径
|
||
# 格式: base_model.model.model.layers.X.self_attn.q_proj.lora_A.weight
|
||
parts = a_key.split(".")
|
||
# 去掉 base_model.model 前缀和 .lora_A.weight 后缀
|
||
# 找到实际模块路径
|
||
try:
|
||
lora_a_idx = parts.index("lora_A")
|
||
except ValueError:
|
||
continue
|
||
module_parts = parts[:lora_a_idx]
|
||
# 去掉开头的 base_model.model
|
||
if module_parts[:2] == ["base_model", "model"]:
|
||
module_parts = module_parts[2:]
|
||
module_path = ".".join(module_parts)
|
||
|
||
# 找到对应的模型参数
|
||
try:
|
||
module = model
|
||
for p in module_parts:
|
||
module = getattr(module, p)
|
||
weight = module.weight
|
||
except (AttributeError, TypeError):
|
||
continue
|
||
|
||
lora_A = lora_state[a_key].to(torch.float32)
|
||
lora_B = lora_state[b_key].to(torch.float32)
|
||
delta = (lora_B @ lora_A) * scaling
|
||
|
||
with torch.no_grad():
|
||
weight.data += delta.to(weight.dtype)
|
||
merged += 1
|
||
|
||
print(f" ✅ 成功合并 {merged} 个 LoRA 层")
|
||
return model
|
||
|
||
|
||
def load_model():
|
||
"""加载基座模型 + 手动合并LoRA权重"""
|
||
import torch
|
||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||
|
||
print("=" * 60)
|
||
print(" Angel Embrace - 16质点心理辅导 AI")
|
||
print(" CPU推理模式启动中...")
|
||
print("=" * 60)
|
||
|
||
# 确定模型加载路径
|
||
if os.path.isdir(LOCAL_MODEL_DIR) and os.path.isfile(os.path.join(LOCAL_MODEL_DIR, "config.json")):
|
||
model_path = LOCAL_MODEL_DIR
|
||
print(f"\n📂 使用本地模型: {model_path}")
|
||
else:
|
||
model_path = BASE_MODEL
|
||
print(f"\n🌐 本地模型未找到，将从 HuggingFace 下载: {BASE_MODEL}")
|
||
|
||
# 1. 加载 Tokenizer
|
||
print("\n📥 [1/4] 正在加载 Tokenizer...")
|
||
tokenizer = AutoTokenizer.from_pretrained(
|
||
model_path,
|
||
trust_remote_code=True,
|
||
)
|
||
print(" ✅ Tokenizer 加载完成")
|
||
|
||
# 2. 加载基座模型 (CPU模式)
|
||
print(f"\n📥 [2/4] 正在加载基座模型（约72GB，需要几分钟）...")
|
||
model = AutoModelForCausalLM.from_pretrained(
|
||
model_path,
|
||
dtype=torch.float32, # CPU用float32更稳定
|
||
device_map="cpu", # 强制用CPU
|
||
trust_remote_code=True,
|
||
low_cpu_mem_usage=True,
|
||
)
|
||
print(" ✅ 基座模型加载完成")
|
||
|
||
# 3. 手动合并 LoRA 权重（不依赖peft）
|
||
print(f"\n📥 [3/4] 正在合并 LoRA 权重...")
|
||
if os.path.isdir(LORA_PATH) and os.path.isfile(os.path.join(LORA_PATH, "adapter_model.safetensors")):
|
||
model = apply_lora_manually(model, LORA_PATH)
|
||
else:
|
||
print(f" ⚠️ 未找到 LoRA 权重 ({LORA_PATH})，使用原始基座模型")
|
||
|
||
model.eval()
|
||
print(f"\n✅ [4/4] 模型就绪！进入对话模式\n")
|
||
|
||
return model, tokenizer
|
||
|
||
|
||
def chat_loop(model, tokenizer):
|
||
"""主对话循环"""
|
||
import torch
|
||
|
||
messages = [
|
||
{"role": "system", "content": SYSTEM_PROMPT}
|
||
]
|
||
|
||
print("=" * 60)
|
||
print(" 输入消息开始对话 | 输入 clear 清空历史 | 输入 quit 退出")
|
||
print("=" * 60)
|
||
|
||
while True:
|
||
try:
|
||
user_input = input("\n你: ").strip()
|
||
except (EOFError, KeyboardInterrupt):
|
||
print("\n👋 再见！照顾好自己 💙")
|
||
break
|
||
|
||
if not user_input:
|
||
continue
|
||
if user_input.lower() in ("quit", "exit", "q", "退出"):
|
||
print("\n👋 再见！照顾好自己 💙")
|
||
break
|
||
if user_input.lower() in ("clear", "cls", "清空"):
|
||
messages = [{"role": "system", "content": SYSTEM_PROMPT}]
|
||
print(" 🗑️ 对话历史已清空")
|
||
continue
|
||
|
||
messages.append({"role": "user", "content": user_input})
|
||
|
||
text = tokenizer.apply_chat_template(
|
||
messages,
|
||
tokenize=False,
|
||
add_generation_prompt=True,
|
||
)
|
||
|
||
inputs = tokenizer(text, return_tensors="pt").to("cpu")
|
||
|
||
if inputs["input_ids"].shape[1] > CONTEXT_LENGTH:
|
||
print(f" ⚠️ 输入过长({inputs['input_ids'].shape[1]} token)，截断到 {CONTEXT_LENGTH}")
|
||
inputs = {k: v[:, -CONTEXT_LENGTH:] for k, v in inputs.items()}
|
||
|
||
print("\nAngel Embrace: ", end="", flush=True)
|
||
|
||
with torch.no_grad():
|
||
output_ids = model.generate(
|
||
**inputs,
|
||
max_new_tokens=MAX_NEW_TOKENS,
|
||
do_sample=True,
|
||
temperature=0.7,
|
||
top_p=0.9,
|
||
repetition_penalty=1.1,
|
||
pad_token_id=tokenizer.pad_token_id or tokenizer.eos_token_id,
|
||
)
|
||
|
||
new_tokens = output_ids[0][inputs["input_ids"].shape[-1]:]
|
||
response = tokenizer.decode(new_tokens, skip_special_tokens=True)
|
||
print(response)
|
||
|
||
messages.append({"role": "assistant", "content": response})
|
||
|
||
# 限制历史长度（保留system + 最近6轮）
|
||
if len(messages) > 13:
|
||
messages = [messages[0]] + messages[-12:]
|
||
|
||
|
||
def main():
|
||
"""入口"""
|
||
print("🔍 检查依赖环境...")
|
||
check_dependencies()
|
||
|
||
model, tokenizer = load_model()
|
||
chat_loop(model, tokenizer)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|