Upload 2 files

vvvvvv/app.py

@@ -0,0 +1,174 @@
+import gradio as gr
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import json
+import os
+from train import ModelTrainer
+
+class NovelAIApp:
+    def __init__(self):
+        self.model = None
+        self.tokenizer = None
+        self.trainer = None
+        
+        # 加载系统提示词
+        with open('configs/system_prompts.json', 'r', encoding='utf-8') as f:
+            self.system_prompts = json.load(f)
+        
+        # 初始化默认的情境
+        self.current_mood = "暗示"
+        self.mood_history = []      # 情境历史记录
+
+    def load_model(self, model_path):
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            model_path,
+            trust_remote_code=True
+        )
+        self.model = AutoModelForCausalLM.from_pretrained(
+            model_path,
+            trust_remote_code=True,
+            load_in_8bit=True,
+            device_map="auto"
+        )
+
+    def train_model(self, files):
+        if not self.trainer:
+            self.trainer = ModelTrainer(
+                "THUDM/chatglm2-6b",
+                "configs/system_prompts.json"
+            )
+        
+        dataset = self.trainer.prepare_dataset(files)
+        self.trainer.train(dataset)
+        return "训练完成！"
+
+    def generate_text(self, message, history, mood=None):
+        """增强的生成文本方法"""
+        if not self.model:
+            return "请先加载模型！"
+        
+        # 智能情境切换
+        if mood:
+            self.current_mood = mood
+        else:
+            # 根据用户输入自动判断情境
+            self.current_mood = self._detect_mood(message)
+        
+        system_prompt = f"""<|system|>{self.system_prompts['base_prompt']}
+当前情境：{self.current_mood}</|system|>"""
+        
+        # 构建带情感的对话历史
+        full_history = ""
+        for msg in history:
+            full_history += f"<|user|>{msg[0]}</|user|>\n<|assistant|>{msg[1]}</|assistant|>\n"
+        
+        formatted_prompt = f"{system_prompt}\n{full_history}<|user|>{message}</|user|>\n<|assistant|>"
+        
+        # 增加温度动态调整
+        temperature = self._get_dynamic_temperature()
+        
+        inputs = self.tokenizer(formatted_prompt, return_tensors="pt")
+        outputs = self.model.generate(
+            inputs["input_ids"],
+            max_length=1024,
+            temperature=temperature,
+            top_p=0.9,
+            repetition_penalty=1.1,
+            do_sample=True,          # 启用采样
+            num_return_sequences=1,
+            pad_token_id=self.tokenizer.eos_token_id
+        )
+        
+        response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+        response = response.split("<|assistant|>")[-1].strip()
+        
+        # 记录情境历史
+        self.mood_history.append(self.current_mood)
+        return response
+
+    def _detect_mood(self, message):
+        """根据用户输入智能判断情境"""
+        mood_keywords = {
+            "撒娇": ["想你", "抱抱", "亲亲", "摸摸"],
+            "害羞": ["害羞", "不好意思", "脸红"],
+            "粘人": ["陪我", "一起", "不要走"],
+            "暗示": ["热", "难受", "想要", "摸"],
+            "调皮": ["玩", "闹", "捣乱", "逗"],
+            "吃醋": ["别人", "不理我", "生气"]
+        }
+        
+        # 计算关键词匹配度
+        mood_scores = {mood: 0 for mood in mood_keywords}
+        for mood, keywords in mood_keywords.items():
+            for keyword in keywords:
+                if keyword in message:
+                    mood_scores[mood] += 1
+        
+        # 如果没有明显倾向，保持当前情境
+        max_score = max(mood_scores.values())
+        if max_score == 0:
+            return self.current_mood
+            
+        return max(mood_scores.items(), key=lambda x: x[1])[0]
+
+    def _get_dynamic_temperature(self):
+        """根据情境动态调整生成温度"""
+        temperature_map = {
+            "撒娇": 0.8,
+            "害羞": 0.6,
+            "粘人": 0.7,
+            "暗示": 0.9,
+            "调皮": 0.85,
+            "吃醋": 0.75
+        }
+        return temperature_map.get(self.current_mood, 0.7)
+
+    def create_interface(self):
+        """优化的界面创建方法"""
+        with gr.Blocks() as interface:
+            gr.Markdown("# 猫娘对话助手")
+            
+            with gr.Tab("模型训练"):
+                with gr.Row():
+                    file_output = gr.File(
+                        file_count="multiple",
+                        label="上传小说文本文件"
+                    )
+                    train_button = gr.Button("开始训练")
+                
+                train_output = gr.Textbox(label="训练状态")
+                
+            with gr.Tab("对话"):
+                with gr.Row():
+                    mood_selector = gr.Dropdown(
+                        choices=["撒娇", "害羞", "粘人", "暗示", "调皮", "吃醋"],
+                        value=self.current_mood,
+                        label="选择当前情境"
+                    )
+                    
+                chatbot = gr.ChatInterface(
+                    fn=lambda msg, history: self.generate_text(msg, history, mood_selector.value),
+                    title="与猫娘对话",
+                    description="来和可爱的猫娘聊天吧~",
+                    theme="soft",
+                    examples=[
+                        "今天好想你呀~",
+                        "主人在做什么呢？",
+                        "要不要一起玩？",
+                        "人家身体有点奇怪...",
+                        "主人不要理别人啦！"
+                    ],
+                    cache_examples=False
+                )
+
+        return interface
+
+# 创建应用实例
+app = NovelAIApp()
+interface = app.create_interface()
+
+# 修改 launch 参数
+interface.launch(
+    server_name="0.0.0.0",  # 允许外部访问
+    share=True,             # 创建公共链接
+    ssl_verify=False        # 禁用 SSL 验证
+) 

vvvvvv/train.py

@@ -0,0 +1,339 @@
+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
+
+class ModelTrainer:
+    def __init__(self, model_id, system_prompts_path):
+        self.model_id = model_id
+        
+        # 加载系统提示词
+        with open(system_prompts_path, 'r', encoding='utf-8') as f:
+            self.system_prompts = json.load(f)
+        
+        # 初始化tokenizer和model - 添加trust_remote_code=True
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            model_id,
+            trust_remote_code=True  # 添加此参数
+        )
+        self.model = AutoModelForCausalLM.from_pretrained(
+            model_id,
+            trust_remote_code=True,  # 添加此参数
+            low_cpu_mem_usage=True,
+            torch_dtype='float32'
+        )
+        
+        # 使用更轻量的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
+
+        def _create_mixed_mood_sample(self, chunk):
+            """创建混合情境的训练样本"""
+            moods = random.sample(list(dialogue_contexts.keys()), 2)
+            mood1, mood2 = moods
+            
+            # 处理两种情境的文本
+            processed_chunk1 = self._process_text_style(chunk, mood1)
+            processed_chunk2 = self._process_text_style(chunk, mood2)
+            
+            # 创建混合情境对话
+            conversation = f"""<|system|>{base_system_prompt}
+当前情境：{mood1}转{mood2}</|system|>
+<|user|>{random.choice(dialogue_contexts[mood1])['question']}</|user|>
+<|assistant|>{processed_chunk1}</|assistant|>
+<|user|>{random.choice(dialogue_contexts[mood2])['question']}</|user|>
+<|assistant|>{processed_chunk2}</|assistant|>"""
+            
+            return conversation
+
+        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
+
+                    # 在现有的训练数据生成循环中添加混合情境样本
+                    if random.random() < 0.2:  # 20%的概率生成混合情境样本
+                        mixed_conversation = self._create_mixed_mood_sample(chunk)
+                        dataset.append({"text": mixed_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)