""" 多聊天室表达风格学习系统 支持为每个chat_id维护独立的表达模型，学习从up_content到style的映射 """ import os import pickle import traceback from typing import Dict, List, Optional, Tuple from collections import defaultdict import asyncio from src.common.logger import get_logger from .expressor_model.model import ExpressorModel logger = get_logger("style_learner") class StyleLearner: """ 单个聊天室的表达风格学习器 学习从up_content到style的映射关系 支持动态管理风格集合（无数量上限） """ def __init__(self, chat_id: str, model_config: Optional[Dict] = None): self.chat_id = chat_id self.model_config = model_config or { "alpha": 0.5, "beta": 0.5, "gamma": 0.99, # 衰减因子，支持遗忘 "vocab_size": 200000, "use_jieba": True, } # 初始化表达模型 self.expressor = ExpressorModel(**self.model_config) # 动态风格管理 self.style_to_id: Dict[str, str] = {} # style文本 -> style_id self.id_to_style: Dict[str, str] = {} # style_id -> style文本 self.id_to_situation: Dict[str, str] = {} # style_id -> situation文本 self.next_style_id = 0 # 下一个可用的style_id # 学习统计 self.learning_stats = { "total_samples": 0, "style_counts": defaultdict(int), "last_update": None, "style_usage_frequency": defaultdict(int), # 风格使用频率 } def add_style(self, style: str, situation: str = None) -> bool: """ 动态添加一个新的风格 Args: style: 风格文本 situation: 对应的situation文本（可选） Returns: bool: 添加是否成功 """ try: # 检查是否已存在 if style in self.style_to_id: logger.debug(f"[{self.chat_id}] 风格 '{style}' 已存在") return True # 生成新的style_id style_id = f"style_{self.next_style_id}" self.next_style_id += 1 # 添加到映射 self.style_to_id[style] = style_id self.id_to_style[style_id] = style if situation: self.id_to_situation[style_id] = situation # 添加到expressor模型 self.expressor.add_candidate(style_id, style, situation) logger.info( f"[{self.chat_id}] 已添加风格: '{style}' (ID: {style_id})" + (f", situation: '{situation}'" if situation else "") ) return True except Exception as e: logger.error(f"[{self.chat_id}] 添加风格失败: {e}") return False def remove_style(self, style: str) -> bool: """ 删除一个风格 Args: style: 要删除的风格文本 Returns: bool: 删除是否成功 """ try: if style not in self.style_to_id: logger.warning(f"[{self.chat_id}] 风格 '{style}' 不存在") return False style_id = self.style_to_id[style] # 从映射中删除 del self.style_to_id[style] del self.id_to_style[style_id] if style_id in self.id_to_situation: del self.id_to_situation[style_id] # 从expressor模型中删除（通过重新构建） self._rebuild_expressor() logger.info(f"[{self.chat_id}] 已删除风格: '{style}' (ID: {style_id})") return True except Exception as e: logger.error(f"[{self.chat_id}] 删除风格失败: {e}") return False def update_style(self, old_style: str, new_style: str) -> bool: """ 更新一个风格 Args: old_style: 原风格文本 new_style: 新风格文本 Returns: bool: 更新是否成功 """ try: if old_style not in self.style_to_id: logger.warning(f"[{self.chat_id}] 原风格 '{old_style}' 不存在") return False if new_style in self.style_to_id and new_style != old_style: logger.warning(f"[{self.chat_id}] 新风格 '{new_style}' 已存在") return False style_id = self.style_to_id[old_style] # 更新映射 del self.style_to_id[old_style] self.style_to_id[new_style] = style_id self.id_to_style[style_id] = new_style # 更新expressor模型（保留原有的situation） situation = self.id_to_situation.get(style_id) self.expressor.add_candidate(style_id, new_style, situation) logger.info(f"[{self.chat_id}] 已更新风格: '{old_style}' -> '{new_style}'") return True except Exception as e: logger.error(f"[{self.chat_id}] 更新风格失败: {e}") return False def add_styles_batch(self, styles: List[str], situations: List[str] = None) -> int: """ 批量添加风格 Args: styles: 风格文本列表 situations: 对应的situation文本列表（可选） Returns: int: 成功添加的数量 """ success_count = 0 for i, style in enumerate(styles): situation = situations[i] if situations and i < len(situations) else None if self.add_style(style, situation): success_count += 1 logger.info(f"[{self.chat_id}] 批量添加风格: {success_count}/{len(styles)} 成功") return success_count def get_all_styles(self) -> List[str]: """获取所有已注册的风格""" return list(self.style_to_id.keys()) def get_style_count(self) -> int: """获取当前风格数量""" return len(self.style_to_id) def get_situation(self, style: str) -> Optional[str]: """ 获取风格对应的situation Args: style: 风格文本 Returns: Optional[str]: 对应的situation，如果不存在则返回None """ if style not in self.style_to_id: return None style_id = self.style_to_id[style] return self.id_to_situation.get(style_id) def get_style_info(self, style: str) -> Tuple[Optional[str], Optional[str]]: """ 获取风格的完整信息 Args: style: 风格文本 Returns: Tuple[Optional[str], Optional[str]]: (style_id, situation) """ if style not in self.style_to_id: return None, None style_id = self.style_to_id[style] situation = self.id_to_situation.get(style_id) return style_id, situation def get_all_style_info(self) -> Dict[str, Tuple[str, Optional[str]]]: """ 获取所有风格的完整信息 Returns: Dict[str, Tuple[str, Optional[str]]]: {style: (style_id, situation)} """ result = {} for style, style_id in self.style_to_id.items(): situation = self.id_to_situation.get(style_id) result[style] = (style_id, situation) return result def _rebuild_expressor(self): """重新构建expressor模型（删除风格后使用）""" try: # 重新创建expressor self.expressor = ExpressorModel(**self.model_config) # 重新添加所有风格和situation for style_id, style_text in self.id_to_style.items(): situation = self.id_to_situation.get(style_id) self.expressor.add_candidate(style_id, style_text, situation) logger.debug(f"[{self.chat_id}] 已重新构建expressor模型") except Exception as e: logger.error(f"[{self.chat_id}] 重新构建expressor失败: {e}") def learn_mapping(self, up_content: str, style: str) -> bool: """ 学习一个up_content到style的映射 如果style不存在，会自动添加 Args: up_content: 输入内容 style: 对应的style文本 Returns: bool: 学习是否成功 """ try: # 如果style不存在，先添加它 if style not in self.style_to_id: if not self.add_style(style): logger.warning(f"[{self.chat_id}] 无法添加风格 '{style}'，学习失败") return False # 获取style_id style_id = self.style_to_id[style] # 使用正反馈学习 self.expressor.update_positive(up_content, style_id) # 更新统计 self.learning_stats["total_samples"] += 1 self.learning_stats["style_counts"][style_id] += 1 self.learning_stats["style_usage_frequency"][style] += 1 self.learning_stats["last_update"] = asyncio.get_event_loop().time() logger.debug(f"[{self.chat_id}] 学习映射: '{up_content}' -> '{style}'") return True except Exception as e: logger.error(f"[{self.chat_id}] 学习映射失败: {e}") traceback.print_exc() return False def predict_style(self, up_content: str, top_k: int = 5) -> Tuple[Optional[str], Dict[str, float]]: """ 根据up_content预测最合适的style Args: up_content: 输入内容 top_k: 返回前k个候选 Returns: Tuple[最佳style文本, 所有候选的分数] """ try: best_style_id, scores = self.expressor.predict(up_content, k=top_k) if best_style_id is None: return None, {} # 将style_id转换为style文本 best_style = self.id_to_style.get(best_style_id) # 转换所有分数 style_scores = {} for sid, score in scores.items(): style_text = self.id_to_style.get(sid) if style_text: style_scores[style_text] = score return best_style, style_scores except Exception as e: logger.error(f"[{self.chat_id}] 预测style失败: {e}") traceback.print_exc() return None, {} def decay_learning(self, factor: Optional[float] = None) -> None: """ 对学习到的知识进行衰减（遗忘） Args: factor: 衰减因子，None则使用配置中的gamma """ self.expressor.decay(factor) logger.debug(f"[{self.chat_id}] 执行知识衰减") def get_stats(self) -> Dict: """获取学习统计信息""" return { "chat_id": self.chat_id, "total_samples": self.learning_stats["total_samples"], "style_count": len(self.style_to_id), "style_counts": dict(self.learning_stats["style_counts"]), "style_usage_frequency": dict(self.learning_stats["style_usage_frequency"]), "last_update": self.learning_stats["last_update"], "all_styles": list(self.style_to_id.keys()), } def save(self, base_path: str) -> bool: """ 保存模型到文件 Args: base_path: 基础路径，实际文件为 {base_path}/{chat_id}_style_model.pkl """ try: os.makedirs(base_path, exist_ok=True) file_path = os.path.join(base_path, f"{self.chat_id}_style_model.pkl") # 保存模型和统计信息 save_data = { "model_config": self.model_config, "style_to_id": self.style_to_id, "id_to_style": self.id_to_style, "id_to_situation": self.id_to_situation, "next_style_id": self.next_style_id, "learning_stats": self.learning_stats, } # 先保存expressor模型 expressor_path = os.path.join(base_path, f"{self.chat_id}_expressor.pkl") self.expressor.save(expressor_path) # 保存其他数据 with open(file_path, "wb") as f: pickle.dump(save_data, f) logger.info(f"[{self.chat_id}] 模型已保存到 {file_path}") return True except Exception as e: logger.error(f"[{self.chat_id}] 保存模型失败: {e}") return False def load(self, base_path: str) -> bool: """ 从文件加载模型 Args: base_path: 基础路径 """ try: file_path = os.path.join(base_path, f"{self.chat_id}_style_model.pkl") expressor_path = os.path.join(base_path, f"{self.chat_id}_expressor.pkl") if not os.path.exists(file_path) or not os.path.exists(expressor_path): logger.warning(f"[{self.chat_id}] 模型文件不存在，将使用默认配置") return False # 加载其他数据 with open(file_path, "rb") as f: save_data = pickle.load(f) # 恢复配置和状态 self.model_config = save_data["model_config"] self.style_to_id = save_data["style_to_id"] self.id_to_style = save_data["id_to_style"] self.id_to_situation = save_data.get("id_to_situation", {}) # 兼容旧版本 self.next_style_id = save_data["next_style_id"] self.learning_stats = save_data["learning_stats"] # 重新创建expressor并加载 self.expressor = ExpressorModel(**self.model_config) self.expressor.load(expressor_path) logger.info(f"[{self.chat_id}] 模型已从 {file_path} 加载") return True except Exception as e: logger.error(f"[{self.chat_id}] 加载模型失败: {e}") return False class StyleLearnerManager: """ 多聊天室表达风格学习管理器 为每个chat_id维护独立的StyleLearner实例 每个chat_id可以动态管理自己的风格集合（无数量上限） """ def __init__(self, model_save_path: str = "data/style_models"): self.model_save_path = model_save_path self.learners: Dict[str, StyleLearner] = {} # 自动保存配置 self.auto_save_interval = 300 # 5分钟 self._auto_save_task: Optional[asyncio.Task] = None logger.info("StyleLearnerManager 已初始化") def get_learner(self, chat_id: str, model_config: Optional[Dict] = None) -> StyleLearner: """ 获取或创建指定chat_id的学习器 Args: chat_id: 聊天室ID model_config: 模型配置，None则使用默认配置 Returns: StyleLearner实例 """ if chat_id not in self.learners: # 创建新的学习器 learner = StyleLearner(chat_id, model_config) # 尝试加载已保存的模型 learner.load(self.model_save_path) self.learners[chat_id] = learner logger.info(f"为 chat_id={chat_id} 创建新的StyleLearner") return self.learners[chat_id] def add_style(self, chat_id: str, style: str) -> bool: """ 为指定chat_id添加风格 Args: chat_id: 聊天室ID style: 风格文本 Returns: bool: 添加是否成功 """ learner = self.get_learner(chat_id) return learner.add_style(style) def remove_style(self, chat_id: str, style: str) -> bool: """ 为指定chat_id删除风格 Args: chat_id: 聊天室ID style: 风格文本 Returns: bool: 删除是否成功 """ learner = self.get_learner(chat_id) return learner.remove_style(style) def update_style(self, chat_id: str, old_style: str, new_style: str) -> bool: """ 为指定chat_id更新风格 Args: chat_id: 聊天室ID old_style: 原风格文本 new_style: 新风格文本 Returns: bool: 更新是否成功 """ learner = self.get_learner(chat_id) return learner.update_style(old_style, new_style) def get_chat_styles(self, chat_id: str) -> List[str]: """ 获取指定chat_id的所有风格 Args: chat_id: 聊天室ID Returns: List[str]: 风格列表 """ learner = self.get_learner(chat_id) return learner.get_all_styles() def learn_mapping(self, chat_id: str, up_content: str, style: str) -> bool: """ 学习一个映射关系 Args: chat_id: 聊天室ID up_content: 输入内容 style: 对应的style Returns: bool: 学习是否成功 """ learner = self.get_learner(chat_id) return learner.learn_mapping(up_content, style) def predict_style(self, chat_id: str, up_content: str, top_k: int = 5) -> Tuple[Optional[str], Dict[str, float]]: """ 预测最合适的style Args: chat_id: 聊天室ID up_content: 输入内容 top_k: 返回前k个候选 Returns: Tuple[最佳style, 所有候选分数] """ learner = self.get_learner(chat_id) return learner.predict_style(up_content, top_k) def decay_all_learners(self, factor: Optional[float] = None) -> None: """ 对所有学习器执行衰减 Args: factor: 衰减因子 """ for learner in self.learners.values(): learner.decay_learning(factor) logger.info("已对所有学习器执行衰减") def get_all_stats(self) -> Dict[str, Dict]: """获取所有学习器的统计信息""" return {chat_id: learner.get_stats() for chat_id, learner in self.learners.items()} def save_all_models(self) -> bool: """保存所有模型""" success_count = 0 for learner in self.learners.values(): if learner.save(self.model_save_path): success_count += 1 logger.info(f"已保存 {success_count}/{len(self.learners)} 个模型") return success_count == len(self.learners) def load_all_models(self) -> int: """加载所有已保存的模型""" if not os.path.exists(self.model_save_path): return 0 loaded_count = 0 for filename in os.listdir(self.model_save_path): if filename.endswith("_style_model.pkl"): chat_id = filename.replace("_style_model.pkl", "") learner = StyleLearner(chat_id) if learner.load(self.model_save_path): self.learners[chat_id] = learner loaded_count += 1 logger.info(f"已加载 {loaded_count} 个模型") return loaded_count async def start_auto_save(self) -> None: """启动自动保存任务""" if self._auto_save_task is None or self._auto_save_task.done(): self._auto_save_task = asyncio.create_task(self._auto_save_loop()) logger.info("已启动自动保存任务") async def stop_auto_save(self) -> None: """停止自动保存任务""" if self._auto_save_task and not self._auto_save_task.done(): self._auto_save_task.cancel() try: await self._auto_save_task except asyncio.CancelledError: pass logger.info("已停止自动保存任务") async def _auto_save_loop(self) -> None: """自动保存循环""" while True: try: await asyncio.sleep(self.auto_save_interval) self.save_all_models() except asyncio.CancelledError: break except Exception as e: logger.error(f"自动保存失败: {e}") # 全局管理器实例 style_learner_manager = StyleLearnerManager()