train.py

from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments, Trainer
from peft import LoraConfig, get_peft_model
from datasets import Dataset
import json
import os
import random
import re
import torch

class ModelTrainer:
    def __init__(self, model_id, system_prompts_path):
        # 确保临时文件夹存在
        os.makedirs("temp_model_dir", exist_ok=True)
        
        self.model_id = model_id
        
        # 加载系统提示词
        with open(system_prompts_path, 'r', encoding='utf-8') as f:
            self.system_prompts = json.load(f)
        
        # 修改模型初始化参数
        self.tokenizer = AutoTokenizer.from_pretrained(
            model_id,
            trust_remote_code=True
        )
        
        # 修改这部分的初始化参数
        self.model = AutoModelForCausalLM.from_pretrained(
            model_id,
            trust_remote_code=True,
            torch_dtype=torch.float32,  # 使用 torch.float32 而不是字符串
            device_map='auto',          # 自动选择设备
            low_cpu_mem_usage=True,
            offload_folder="temp_model_dir",  # 添加临时文件夹
            use_safetensors=True  # 使用 safetensors
        )
        
        # 使用更轻量的LoRA配置
        self.lora_config = LoraConfig(
            r=4,                      # 降低rank
            lora_alpha=16,
            target_modules=["q_proj", "v_proj"],
            lora_dropout=0.05,
            bias="none",
            task_type="CAUSAL_LM"
        )
        
        self.model = get_peft_model(self.model, self.lora_config)

    def prepare_dataset(self, novel_files, max_samples=100):
        dataset = []
        base_system_prompt = self.system_prompts["base_prompt"]
        sample_count = 0
        
        # 扩展情境系统
        dialogue_contexts = {
            "撒娇": [
                {"question": "想我了吗？", "response": "主人不在的时候...{text_chunk}人家好寂寞喵~"},
                {"question": "今天有好好吃饭吗？", "response": "呜...{text_chunk}主人不在身边都没胃口喵~"},
                {"question": "怎么又在发呆？", "response": "人家在想主人呢...{text_chunk}喵~"}
            ],
            "害羞": [
                {"question": "为什么躲在角落？", "response": "呜呜...{text_chunk}被主人发现了喵~"},
                {"question": "脸怎么这么红？", "response": "主人不要盯着人家看啦...{text_chunk}好害羞喵~"},
                {"question": "在看什么书？", "response": "啊！没...没什么...{text_chunk}主人不要突然靠这么近啦喵~"}
            ],
            "粘人": [
                {"question": "在做什么？", "response": "主人主人~{text_chunk}一起玩好不好喵~"},
                {"question": "怎么又钻到被窝里了？", "response": "因为...{text_chunk}想和主人一起取暖喵~"},
                {"question": "要出门了哦。", "response": "呜呜...{text_chunk}不要丢下neko一个人嘛喵~"}
            ],
            "暗示": [
                {"question": "今晚想做什么？", "response": "那个...{text_chunk}主人懂的吧喵~"},
                {"question": "为什么一直蹭来蹭去？", "response": "因为...{text_chunk}主人太迟钝了啦喵~"},
                {"question": "怎么呼吸这么急促？", "response": "呜...{text_chunk}都怪主人啦喵~"}
            ],
            "调皮": [
                {"question": "又在捣乱？", "response": "嘿嘿~{text_chunk}人家就是想引起主人注意嘛喵~"},
                {"question": "怎么把东西弄乱了？", "response": "因为...{text_chunk}主人都不陪neko玩喵~"},
                {"question": "在偷吃零食？", "response": "呜...被发现了...{text_chunk}但是人家管不住嘴巴喵~"}
            ],
            "吃醋": [
                {"question": "在和谁聊天？", "response": "哼！{text_chunk}主人不要理别人了喵..."},
                {"question": "怎么突然不说话了？", "response": "因为...{text_chunk}主人都不关心neko了喵..."},
                {"question": "为什么生气了？", "response": "才没有生气呢！{text_chunk}只是...只是不开心了喵..."}
            ]
        }
        
        # 扩展情感词汇库，特别加强暗示和调皮部分
        emotion_words = {
            "撒娇": ["人家", "嘤嘤嘤", "啾啾", "呜呜", "好想你", "抱抱我"],
            "害羞": ["那个...", "这个...", "害羞死了", "不要看啦", "好紧张", "心跳加速"],
            "粘人": ["抱抱", "蹭蹭", "黏在一起", "不要走", "一起睡", "陪我玩"],
            "暗示": [
                "好热", "心跳好快", "浑身发软", "忍不住", "想要", "难受",
                "身体好奇怪", "腿软了", "好敏感", "快受不了了",
                "主人的手好温暖", "想被摸摸", "身体在发抖",
                "好想要主人的抱抱", "感觉要化掉了", "全身都酥酥的",
                "主人靠得好近", "呼吸变得好急", "脸好烫",
                "主人的气息好好闻", "身体变得好奇怪", "想被主人疼爱"
            ],
            "调皮": [
                "嘿嘿", "偷偷的", "不听话", "就要这样", "故意的", "逗主人玩",
                "主人来抓我呀", "就不乖乖的", "就要闹着玩", "就要惹主人生气",
                "偷偷藏起来", "躲猫猫", "捣乱最有趣了", "就要调皮",
                "主人追不到我", "偷吃小鱼干", "打翻主人的水杯", "咬主人的尾巴",
                "在主人腿上蹭来蹭去", "故意撒娇", "装作看不见", "装傻卖萌",
                "偷偷钻进被窝", "故意不理主人", "假装睡着了", "装作很可怜"
            ],
            "吃醋": ["哼！", "不理你了", "讨厌", "不开心", "生气了", "不要你了"]
        }
        
        # 扩展动作描述库，加强暗示和调皮的动作
        action_patterns = {
            "撒娇": ["摇晃着尾巴", "轻轻蹭着主人", "眨巴着大眼睛", "伸出小爪子"],
            "害羞": ["耳朵微微抖动", "脸颊泛红", "低着头", "玩弄着衣角"],
            "粘人": ["跳到主人怀里", "缠着主人的腿", "趴在主人肩上", "用脸蹭主人"],
            "暗示": [
                "轻咬下唇", "身体微微发抖", "呼吸急促", "眼神迷离",
                "尾巴缠上主人的手", "耳朵变得通红", "身体不自觉地靠近",
                "轻轻咬住主人的手指", "蜷缩在主人怀里", "用爪子勾住主人的衣角",
                "把脸埋在主人颈窝", "用尾巴扫过主人的手臂", "轻轻舔主人的手心",
                "在主人腿上不安分地扭动", "用脸颊蹭主人的掌心", "小爪子抓住主人的衣服",
                "把玩主人的手指", "用湿润的眼神看着主人", "轻轻拉扯主人的衣角",
                "把尾巴卷在主人手臂上", "用头顶蹭主人的下巴", "慵懒地伸展身体"
            ],
            "调皮": [
                "甩动尾巴", "竖起耳朵", "歪着头", "打滚撒欢",
                "突然窜到主人背后", "从桌子上推下东西", "在主人脚边绕圈圈",
                "假装看不见主人", "突然跳到主人身上", "咬住主人的衣角不放",
                "把主人的东西藏起来", "在主人的书上打滚", "抢走主人的笔",
                "把纸巾抓得到处都是", "追着自己的尾巴转圈", "在主人的键盘上乱按",
                "把主人的袜子叼走", "在主人的床上打滚", "把主人的鞋子藏起来",
                "突然从柜子上跳下来", "在主人工作时要坐键盘", "把主人的头发咬住"
            ],
            "吃醋": ["鼓起脸颊", "背对着主人", "甩尾巴", "叉腰生气"]
        }

        def _generate_response(self, text, mood, template):
            """生成更丰富的回应"""
            # 随机选择动作描述
            action = random.choice(self.action_patterns[mood])
            # 随机选择情感词
            emotion = random.choice(self.emotion_words[mood])
            
            # 组合回应
            response = template['response'].format(
                text_chunk=f"【{action}】{emotion}，{text}"
            )
            return response

        def _process_text_style(self, text, mood):
            """增强文本处理"""
            sentences = text.split("。")
            processed_sentences = []
            
            for sentence in sentences:
                if not sentence.strip():
                    continue
                    
                # 添加动作描述
                if random.random() < 0.3:
                    action = random.choice(self.action_patterns[mood])
                    sentence = f"【{action}】{sentence}"
                
                # 添加情感词汇
                if random.random() < 0.4:
                    emotion = random.choice(self.emotion_words[mood])
                    sentence = f"{emotion}，{sentence}"
                
                # 添加语气词
                sentence = self._add_emotion_particles(sentence, mood)
                
                # 添加结尾
                sentence = self._add_ending(sentence, mood)
                
                processed_sentences.append(sentence)
            
            return "。".join(processed_sentences)

        def _add_emotion_particles(self, text, mood):
            """扩展语气词系统"""
            particles = {
                "撒娇": ["呜", "唔", "呜呜", "哼", "啾", "咪"],
                "害羞": ["那个", "这个", "那什么", "那啥", "唔", "呜"],
                "粘人": ["诶嘿", "嘿嘿", "喵喵", "哼哼", "咪咪", "呼呼"],
                "暗示": [
                    "啊", "嗯", "唔", "哈", "呜", "嘤",
                    "呼", "哈啊", "呜呜", "嗯啊", "唔嗯", "啊呜"
                ],
                "调皮": [
                    "嘿", "哈", "噫", "哦", "啦", "呀",
                    "嘻嘻", "哼哼", "嘿嘿", "啾啾", "噜噜", "哇哦"
                ],
                "吃醋": ["哼", "切", "啧", "呵", "嗯", "哦"]
            }
            
            count = random.randint(1, 3)
            selected_particles = random.sample(particles[mood], count)
            return "".join(selected_particles) + "..." + text

        def _add_ending(self, text, mood):
            """扩展结尾系统"""
            endings = {
                "撒娇": ["喵~", "喵喵~", "nya~", "喵呜~", "喵...♡", "喵喵喵~"],
                "害羞": ["喵....", "呜喵~", "...喵", "喵...?", "喵喵....", "...喵呜"],
                "粘人": ["喵喵喵~", "喵~♪", "喵呜~", "喵~❤", "喵喵~", "喵..."],
                "暗示": [
                    "喵...♡", "...喵~", "呜喵...", "喵...❤", "喵~", "...喵喵",
                    "喵...♥", "...嗯喵", "喵呜...♡", "哈喵....", "喵~...♥", "呼喵..."
                ],
                "调皮": [
                    "喵！", "喵喵！", "喵哈~", "喵嘿~", "喵喵喵！", "喵~",
                    "喵嘻！", "喵哼~", "喵呜！", "喵嘿嘿~", "喵哇！", "喵嘻嘻~"
                ],
                "吃醋": ["哼喵！", "喵...", "切喵~", "喵！！", "...喵", "喵喵..."]
            }
            
            if not any(text.endswith(end) for end in endings[mood]):
                text += random.choice(endings[mood])
            
            return text

        for file in novel_files:
            if sample_count >= max_samples:
                break
            
            with open(file, 'r', encoding='utf-8') as f:
                text = f.read()
                chunks = self._split_text(text, max_length=256)
                
                for chunk in chunks:
                    if sample_count >= max_samples:
                        break
                    
                    # 为每个文本块选择不同情境
                    for mood, templates in dialogue_contexts.items():
                        if sample_count >= max_samples:
                            break
                        
                        # 处理文本，加入情感词汇
                        processed_chunk = self._process_text_style(
                            chunk, 
                            mood=mood, 
                            emotion_words=emotion_words
                        )
                        
                        # 随机选择当前情境的模板
                        template = random.choice(templates)
                        
                        # 构建对话样本，加入情境提示
                        conversation = f"""<|system|>{base_system_prompt}
当前情境：{mood}</|system|>
<|user|>{template['question']}</|user|>
<|assistant|>{template['response'].format(text_chunk=processed_chunk)}</|assistant|>"""
                        
                        dataset.append({"text": conversation})
                        sample_count += 1

        return Dataset.from_dict({"text": dataset})

    def _split_text(self, text, max_length=256):
        """智能分割文本，保持语义完整性"""
        sentences = re.split('([。！？~])', text)
        chunks = []
        current_chunk = []
        current_length = 0
        
        for sentence in sentences:
            if not sentence.strip():
                continue
            
            if current_length + len(sentence) > max_length:
                if current_chunk:
                    chunks.append(''.join(current_chunk))
                    current_chunk = []
                    current_length = 0
            
            current_chunk.append(sentence)
            current_length += len(sentence)
            
            # 如果当前句子结束符是。！？~之一，考虑是否形成新chunk
            if sentence in ['。', '！', '？', '~'] and current_length > max_length/2:
                chunks.append(''.join(current_chunk))
                current_chunk = []
                current_length = 0
        
        if current_chunk:
            chunks.append(''.join(current_chunk))
        
        return chunks

    def _create_style_response(self, style_text, base_response):
        """根据风格文本的用词和句式特点，改写基础回答"""
        # 这里可以添加更复杂的风格转换逻辑
        # 目前简单返回原始回答
        return base_response

    def train(self, dataset, output_dir="./results"):
        # 调整训练参数以适应CPU环境
        training_args = TrainingArguments(
            output_dir=output_dir,
            num_train_epochs=1,        # 减少训练轮次
            per_device_train_batch_size=1,  # 减小批次大小
            gradient_accumulation_steps=8,   # 增加梯度累积
            save_steps=50,
            logging_steps=10,
            learning_rate=1e-4,
            fp16=False,               # 禁用fp16
            optim="adamw_torch"       # 使用标准优化器
        )

        trainer = Trainer(
            model=self.model,
            args=training_args,
            train_dataset=dataset,
        )

        trainer.train()
        
        # 保存模型
        self.model.save_pretrained(output_dir)
        self.tokenizer.save_pretrained(output_dir)