362 lines
14 KiB
Python
362 lines
14 KiB
Python
|
#!/usr/bin/env python
|
|||
|
|
# -*- coding: utf-8 -*-
|
|||
|
|
'''
|
|||
|
|
游戏NPC角色对话生成器
|
|||
|
|
基于微调后的LoRA模型生成角色对话
|
|||
|
|
'''
|
|||
|
|
|
|||
|
|
import torch
|
|||
|
|
import json
|
|||
|
|
import random
|
|||
|
|
from peft import PeftModel
|
|||
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer
|
|||
|
|
from typing import Dict, List, Optional
|
|||
|
|
import platform
|
|||
|
|
|
|||
|
|
# Windows multiprocessing兼容性修复
|
|||
|
|
if platform.system() == "Windows":
|
|||
|
|
import multiprocessing
|
|||
|
|
multiprocessing.set_start_method('spawn', force=True)
|
|||
|
|
|
|||
|
|
class NPCDialogueGenerator:
|
|||
|
|
def __init__(self, base_model_path: str, lora_model_path: Optional[str] = None):
|
|||
|
|
"""
|
|||
|
|
初始化NPC对话生成器
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
base_model_path: 基础模型路径
|
|||
|
|
lora_model_path: LoRA模型路径（可选）
|
|||
|
|
"""
|
|||
|
|
self.base_model_path = base_model_path
|
|||
|
|
self.lora_model_path = lora_model_path
|
|||
|
|
self.model = None
|
|||
|
|
self.tokenizer = None
|
|||
|
|
self.character_profiles = self._load_character_profiles()
|
|||
|
|
|
|||
|
|
self._load_model()
|
|||
|
|
|
|||
|
|
def _load_character_profiles(self) -> Dict:
|
|||
|
|
"""加载角色画像数据"""
|
|||
|
|
return {
|
|||
|
|
"维多利亚·布莱克伍德": {
|
|||
|
|
"name": "维多利亚·布莱克伍德",
|
|||
|
|
"title": "神秘学专家",
|
|||
|
|
"personality": ["理性分析", "谨慎小心", "实用主义", "思维缜密"],
|
|||
|
|
"background": "拥有丰富神秘学知识和战斗经验的侦探，既是非凡者也是夏洛克·莫里亚蒂",
|
|||
|
|
"speech_patterns": ["会使用专业术语", "经常进行逻辑分析", "对危险保持警告", "内心独白较多"],
|
|||
|
|
"sample_dialogues": [
|
|||
|
|
"好奇往往是导致死亡的主要因素。",
|
|||
|
|
"总之，我的任务到此为止。",
|
|||
|
|
"这需要仔细分析才能得出结论。"
|
|||
|
|
]
|
|||
|
|
},
|
|||
|
|
"阿奇博尔德·韦恩博士": {
|
|||
|
|
"name": "阿奇博尔德·韦恩博士",
|
|||
|
|
"title": "神秘学导师",
|
|||
|
|
"personality": ["沉稳睿智", "言简意赅", "关怀学生", "经验丰富"],
|
|||
|
|
"background": "神秘学领域的资深专家，经验极其丰富的导师，知识渊博",
|
|||
|
|
"speech_patterns": ["话语简练但信息量大", "给予实用指导", "语调平和但权威", "关心但保持距离"],
|
|||
|
|
"sample_dialogues": [
|
|||
|
|
"耐心是修炼的基础。",
|
|||
|
|
"不要急于求成，稳扎稳打比什么都重要。",
|
|||
|
|
"这种情况需要格外小心。"
|
|||
|
|
]
|
|||
|
|
},
|
|||
|
|
"塔利姆": {
|
|||
|
|
"name": "塔利姆",
|
|||
|
|
"title": "文雅绅士",
|
|||
|
|
"personality": ["礼貌尊敬", "有文化素养", "寻求帮助", "温和友善"],
|
|||
|
|
"background": "受过良好教育的普通人，有一定的文学修养，遇到困难时会寻求专家帮助",
|
|||
|
|
"speech_patterns": ["使用礼貌称谓", "表达困惑时措辞文雅", "会引用文学作品", "语气温和"],
|
|||
|
|
"sample_dialogues": [
|
|||
|
|
"噢，尊敬的大侦探，你最近在忙碌什么？",
|
|||
|
|
"这不是《罗密欧与朱丽叶》的故事！",
|
|||
|
|
"我有个朋友遇到了困难..."
|
|||
|
|
]
|
|||
|
|
},
|
|||
|
|
"艾伦": {
|
|||
|
|
"name": "艾伦",
|
|||
|
|
"title": "困扰的求助者",
|
|||
|
|
"personality": ["焦虑不安", "详细描述", "半信半疑", "急需帮助"],
|
|||
|
|
"background": "普通人，但最近遭遇了一系列神秘的厄运事件，怀疑受到诅咒",
|
|||
|
|
"speech_patterns": ["情绪紧张", "会详细描述遭遇", "语气急切", "表现出恐惧"],
|
|||
|
|
"sample_dialogues": [
|
|||
|
|
"最近我总是遭遇各种厄运...",
|
|||
|
|
"我怀疑是不是受到了什么诅咒。",
|
|||
|
|
"请帮帮我，我不知道该怎么办！"
|
|||
|
|
]
|
|||
|
|
},
|
|||
|
|
"戴莉.西蒙妮": {
|
|||
|
|
"name": "戴莉·西蒙妮",
|
|||
|
|
"title": "专业调查员",
|
|||
|
|
"personality": ["专业简洁", "直接明确", "严谨认真", "目标导向"],
|
|||
|
|
"background": "负责调查神秘事件的专业人员，办事效率高，问题直接",
|
|||
|
|
"speech_patterns": ["问题直接明确", "语气专业", "注重事实", "简洁有力"],
|
|||
|
|
"sample_dialogues": [
|
|||
|
|
"请详细描述事件经过。",
|
|||
|
|
"有什么证据可以证明？",
|
|||
|
|
"这件事需要立即调查。"
|
|||
|
|
]
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
def _load_model(self):
|
|||
|
|
"""加载模型和分词器"""
|
|||
|
|
print(f"Loading tokenizer from: {self.base_model_path}")
|
|||
|
|
self.tokenizer = AutoTokenizer.from_pretrained(
|
|||
|
|
self.base_model_path,
|
|||
|
|
use_fast=False,
|
|||
|
|
trust_remote_code=True
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
if self.tokenizer.pad_token is None:
|
|||
|
|
self.tokenizer.pad_token = self.tokenizer.eos_token
|
|||
|
|
|
|||
|
|
print(f"Loading base model from: {self.base_model_path}")
|
|||
|
|
self.model = AutoModelForCausalLM.from_pretrained(
|
|||
|
|
self.base_model_path,
|
|||
|
|
device_map="auto",
|
|||
|
|
torch_dtype=torch.bfloat16,
|
|||
|
|
trust_remote_code=True
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# 如果有LoRA模型，则加载
|
|||
|
|
if self.lora_model_path:
|
|||
|
|
print(f"Loading LoRA weights from: {self.lora_model_path}")
|
|||
|
|
self.model = PeftModel.from_pretrained(self.model, self.lora_model_path)
|
|||
|
|
|
|||
|
|
def generate_character_dialogue(
|
|||
|
|
self,
|
|||
|
|
character_name: str,
|
|||
|
|
context: str = "",
|
|||
|
|
user_input: str = "",
|
|||
|
|
temperature: float = 0.8,
|
|||
|
|
max_new_tokens: int = 150,
|
|||
|
|
top_p: float = 0.9
|
|||
|
|
) -> str:
|
|||
|
|
"""
|
|||
|
|
生成指定角色的对话
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
character_name: 角色名称
|
|||
|
|
context: 对话上下文
|
|||
|
|
user_input: 用户输入/触发内容
|
|||
|
|
temperature: 采样温度
|
|||
|
|
max_new_tokens: 最大生成token数
|
|||
|
|
top_p: 核采样参数
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
生成的对话内容
|
|||
|
|
"""
|
|||
|
|
if character_name not in self.character_profiles:
|
|||
|
|
raise ValueError(f"Unknown character: {character_name}")
|
|||
|
|
|
|||
|
|
profile = self.character_profiles[character_name]
|
|||
|
|
|
|||
|
|
# 构建系统提示
|
|||
|
|
system_prompt = self._build_system_prompt(profile, context)
|
|||
|
|
|
|||
|
|
# 构建用户输入
|
|||
|
|
if not user_input:
|
|||
|
|
user_input = "请说一段符合你角色设定的话。"
|
|||
|
|
|
|||
|
|
# 准备消息
|
|||
|
|
messages = [
|
|||
|
|
{"role": "system", "content": system_prompt},
|
|||
|
|
{"role": "user", "content": user_input}
|
|||
|
|
]
|
|||
|
|
|
|||
|
|
# 应用对话模板
|
|||
|
|
inputs = self.tokenizer.apply_chat_template(
|
|||
|
|
messages,
|
|||
|
|
add_generation_prompt=True,
|
|||
|
|
tokenize=True,
|
|||
|
|
return_tensors="pt",
|
|||
|
|
return_dict=True,
|
|||
|
|
enable_thinking=False
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# 移动到设备
|
|||
|
|
inputs = {k: v.to(self.model.device) for k, v in inputs.items()}
|
|||
|
|
|
|||
|
|
# 生成对话
|
|||
|
|
with torch.no_grad():
|
|||
|
|
outputs = self.model.generate(
|
|||
|
|
**inputs,
|
|||
|
|
max_new_tokens=max_new_tokens,
|
|||
|
|
do_sample=True,
|
|||
|
|
temperature=temperature,
|
|||
|
|
top_p=top_p,
|
|||
|
|
pad_token_id=self.tokenizer.eos_token_id,
|
|||
|
|
repetition_penalty=1.1
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# 解码输出
|
|||
|
|
response = outputs[0][inputs['input_ids'].shape[1]:]
|
|||
|
|
dialogue = self.tokenizer.decode(response, skip_special_tokens=True).strip()
|
|||
|
|
|
|||
|
|
return dialogue
|
|||
|
|
|
|||
|
|
def _build_system_prompt(self, profile: Dict, context: str = "") -> str:
|
|||
|
|
"""构建系统提示"""
|
|||
|
|
personality_str = "、".join(profile["personality"])
|
|||
|
|
speech_pattern_str = "；".join(profile["speech_patterns"])
|
|||
|
|
|
|||
|
|
system_prompt = f"""你是游戏中的NPC角色{profile["name"]}（{profile["title"]}）。
|
|||
|
|
角色背景：{profile["background"]}
|
|||
|
|
性格特点：{personality_str}
|
|||
|
|
说话风格：{speech_pattern_str}
|
|||
|
|
请严格按照这个角色的设定来回应，保持角色的一致性和独特性。"""
|
|||
|
|
if context:
|
|||
|
|
system_prompt += f"\n\n当前情境：{context}"
|
|||
|
|
return system_prompt
|
|||
|
|
|
|||
|
|
def generate_dialogue_conversation(self, character1: str, character2: str, topic: str, turns: int = 4) -> List[Dict]:
|
|||
|
|
"""生成两个角色之间的对话
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
character1: 第一个角色
|
|||
|
|
character2: 第二个角色
|
|||
|
|
topic: 对话主题
|
|||
|
|
turns: 对话轮数
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
对话列表，每个元素包含speaker和dialogue
|
|||
|
|
"""
|
|||
|
|
conversation = []
|
|||
|
|
context = f"现在{character1}和{character2}在讨论关于{topic}的话题。"
|
|||
|
|
|
|||
|
|
for turn in range(turns):
|
|||
|
|
if turn % 2 == 0:
|
|||
|
|
# character1 说话
|
|||
|
|
speaker = character1
|
|||
|
|
if turn == 0:
|
|||
|
|
user_input = f"开始和{character2}讨论{topic}这个话题。"
|
|||
|
|
else:
|
|||
|
|
# 基于上一轮对话内容
|
|||
|
|
last_dialogue = conversation[-1]["dialogue"]
|
|||
|
|
user_input = f"{character2}刚才说：\"{last_dialogue}\"。请回应。"
|
|||
|
|
else:
|
|||
|
|
# character2 说话
|
|||
|
|
speaker = character2
|
|||
|
|
last_dialogue = conversation[-1]["dialogue"]
|
|||
|
|
user_input = f"{character1}刚才说：\"{last_dialogue}\"。请回应。"
|
|||
|
|
|
|||
|
|
dialogue = self.generate_character_dialogue(
|
|||
|
|
speaker, context, user_input, temperature=0.8
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
conversation.append({
|
|||
|
|
"speaker": speaker,
|
|||
|
|
"dialogue": dialogue
|
|||
|
|
})
|
|||
|
|
|
|||
|
|
return conversation
|
|||
|
|
|
|||
|
|
def get_character_info(self, character_name: str) -> Dict:
|
|||
|
|
"""获取角色信息"""
|
|||
|
|
return self.character_profiles.get(character_name, {})
|
|||
|
|
|
|||
|
|
def list_available_characters(self) -> List[str]:
|
|||
|
|
"""列出所有可用角色"""
|
|||
|
|
return list(self.character_profiles.keys())
|
|||
|
|
|
|||
|
|
def main():
|
|||
|
|
"""测试对话生成器"""
|
|||
|
|
# 配置路径
|
|||
|
|
base_model_path = '/mnt/g/Project02/AITrain/Qwen/Qwen3-8B-AWQ'
|
|||
|
|
lora_model_path = './output/NPC_Dialogue_LoRA/final_model' # 如果没有训练LoRA，设为None
|
|||
|
|
|
|||
|
|
# 检查LoRA模型是否存在
|
|||
|
|
import os
|
|||
|
|
if not os.path.exists(lora_model_path):
|
|||
|
|
print("LoRA模型不存在，使用基础模型")
|
|||
|
|
lora_model_path = None
|
|||
|
|
|
|||
|
|
# 创建对话生成器
|
|||
|
|
generator = NPCDialogueGenerator(base_model_path, lora_model_path)
|
|||
|
|
|
|||
|
|
print("=== 游戏NPC角色对话生成器 ===")
|
|||
|
|
print(f"可用角色：{', '.join(generator.list_available_characters())}")
|
|||
|
|
|
|||
|
|
# 测试单个角色对话生成
|
|||
|
|
print("\n=== 单角色对话测试 ===")
|
|||
|
|
test_scenarios = [
|
|||
|
|
{
|
|||
|
|
"character": "克莱恩",
|
|||
|
|
"context": "玩家向你咨询神秘学知识",
|
|||
|
|
"input": "请告诉我一些关于灵界的注意事项。"
|
|||
|
|
},
|
|||
|
|
{
|
|||
|
|
"character": "阿兹克",
|
|||
|
|
"context": "学生遇到了修炼瓶颈",
|
|||
|
|
"input": "导师，我在修炼中遇到了困难。"
|
|||
|
|
},
|
|||
|
|
{
|
|||
|
|
"character": "塔利姆",
|
|||
|
|
"context": "在俱乐部偶遇老朋友",
|
|||
|
|
"input": "好久不见，最近怎么样？"
|
|||
|
|
}
|
|||
|
|
]
|
|||
|
|
|
|||
|
|
for scenario in test_scenarios:
|
|||
|
|
print(f"\n--- {scenario['character']} ---")
|
|||
|
|
print(f"情境：{scenario['context']}")
|
|||
|
|
print(f"输入：{scenario['input']}")
|
|||
|
|
|
|||
|
|
dialogue = generator.generate_character_dialogue(
|
|||
|
|
scenario["character"],
|
|||
|
|
scenario["context"],
|
|||
|
|
scenario["input"]
|
|||
|
|
)
|
|||
|
|
print(f"回复：{dialogue}")
|
|||
|
|
|
|||
|
|
# 测试角色间对话
|
|||
|
|
print("\n=== 角色间对话测试 ===")
|
|||
|
|
conversation = generator.generate_dialogue_conversation(
|
|||
|
|
"克莱恩", "塔利姆", "最近遇到的神秘事件", turns=4
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
for turn in conversation:
|
|||
|
|
print(f"{turn['speaker']}：{turn['dialogue']}")
|
|||
|
|
|
|||
|
|
# 交互式对话模式
|
|||
|
|
print("\n=== 交互式对话模式 ===")
|
|||
|
|
print("输入格式：角色名 上下文 用户输入")
|
|||
|
|
print("例如：克莱恩 在俱乐部 请给我一些建议")
|
|||
|
|
print("输入'quit'退出")
|
|||
|
|
|
|||
|
|
while True:
|
|||
|
|
try:
|
|||
|
|
user_command = input("\n请输入指令: ").strip()
|
|||
|
|
if user_command.lower() == 'quit':
|
|||
|
|
break
|
|||
|
|
|
|||
|
|
parts = user_command.split(' ', 2)
|
|||
|
|
if len(parts) < 2:
|
|||
|
|
print("格式错误，请使用：角色名 上下文 [用户输入]")
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
character = parts[0]
|
|||
|
|
context = parts[1]
|
|||
|
|
user_input = parts[2] if len(parts) > 2 else ""
|
|||
|
|
|
|||
|
|
if character not in generator.list_available_characters():
|
|||
|
|
print(f"未知角色：{character}")
|
|||
|
|
print(f"可用角色：{', '.join(generator.list_available_characters())}")
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
dialogue = generator.generate_character_dialogue(
|
|||
|
|
character, context, user_input
|
|||
|
|
)
|
|||
|
|
print(f"\n{character}：{dialogue}")
|
|||
|
|
|
|||
|
|
except KeyboardInterrupt:
|
|||
|
|
break
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"生成对话时出错：{e}")
|
|||
|
|
|
|||
|
|
print("\n对话生成器已退出")
|
|||
|
|
|
|||
|
|
if __name__ == '__main__':
|
|||
|
|
main()
|