core/parametric_memory.py

#!/usr/bin/env python3
"""
虫群v8 — 参数化记忆模型 (Parametric Memory Model)

核心理念：模型即数据库
- 传统方案：对话存数据库，检索靠关键词/向量匹配
- 参数化方案：将记忆编码为模型参数，推理时从参数直接生成

技术路线：
- 基座模型：SwarmModel small (11.8M参数)，CPU友好
- LoRA适配器：基座冻结，记忆通过LoRA增量写入
- 渐进式学习：每积累N条交互，微调LoRA参数
- 灾难性遗忘防治：LoRA复合 + 定期蒸馏

与数据库方案的本质区别：
- 数据库：存储离散文本，检索后拼接到prompt
- 参数化：记忆融入模型参数，生成时自然流露个性化
"""

import copy
import hashlib
import json
import math
import os
_MODEL_DIR = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "training", "models")
import re
import sys
import time
import threading
from collections import deque
from datetime import datetime
from typing import Dict, List, Optional, Tuple

import torch
import torch.nn as nn
import torch.nn.functional as F

# ============================================================
# LoRA层 — 低秩适配，记忆写入的载体
# ============================================================

class LoRALayer(nn.Module):
    """
    LoRA低秩适配层
    将记忆编码为低秩矩阵 ΔW = BA
    原始权重W冻结，只训练B和A
    """
    
    def __init__(self, original_linear: nn.Linear, rank: int = 8, alpha: float = 16.0):
        super().__init__()
        self.original = original_linear
        self.rank = rank
        self.alpha = alpha
        self.scaling = alpha / rank
        
        d_out, d_in = original_linear.weight.shape
        
        # LoRA矩阵: B(d_out x r) * A(r x d_in)
        self.lora_A = nn.Parameter(torch.zeros(rank, d_in))
        self.lora_B = nn.Parameter(torch.zeros(d_out, rank))
        
        # 初始化：A用kaiming，B用零 → 初始时ΔW=0，不影响基座
        nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
        nn.init.zeros_(self.lora_B)
        
        # 冻结原始权重
        self.original.weight.requires_grad = False
        if self.original.bias is not None:
            self.original.bias.requires_grad = False
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # 原始变换 + LoRA增量
        original_out = self.original(x)
        lora_out = (x @ self.lora_A.T @ self.lora_B.T) * self.scaling
        return original_out + lora_out


def apply_lora_to_model(model: nn.Module, rank: int = 8, alpha: float = 16.0,
                        target_modules: List[str] = None) -> nn.Module:
    """
    给模型的所有Linear层应用LoRA
    target_modules: 只对指定名称的层应用，None则全部
    """
    if target_modules is None:
        target_modules = ["qkv", "proj", "net.0", "net.2"]  # 注意力+FFN
    
    lora_layers = []
    for name, module in model.named_modules():
        if isinstance(module, nn.Linear):
            # 检查是否在目标模块中
            should_apply = any(t in name for t in target_modules)
            if should_apply:
                # 找到父模块
                parts = name.split('.')
                parent = model
                for part in parts[:-1]:
                    parent = getattr(parent, part)
                
                # 替换为LoRA层
                lora_layer = LoRALayer(module, rank=rank, alpha=alpha)
                setattr(parent, parts[-1], lora_layer)
                lora_layers.append(name)
    
    return model, lora_layers


def get_lora_params(model: nn.Module) -> List[nn.Parameter]:
    """获取所有LoRA可训练参数"""
    params = []
    for name, param in model.named_parameters():
        if param.requires_grad and ('lora_A' in name or 'lora_B' in name):
            params.append(param)
    return params


def get_lora_state(model: nn.Module) -> Dict[str, torch.Tensor]:
    """导出LoRA参数（用于保存记忆快照）"""
    state = {}
    for name, param in model.named_parameters():
        if 'lora_A' in name or 'lora_B' in name:
            state[name] = param.data.clone()
    return state


def set_lora_state(model: nn.Module, state: Dict[str, torch.Tensor]):
    """加载LoORA参数（用于恢复记忆）"""
    for name, param in model.named_parameters():
        if name in state:
            param.data.copy_(state[name])


# ============================================================
# 记忆编码器 — 将交互数据编码为训练样本
# ============================================================

class MemoryEncoder:
    """
    将用户交互编码为模型可学习的训练样本
    
    格式: [BOS]用户:xxx[SEP]助手:xxx[EOS]
    
    关键设计：
    - 不同类型记忆用不同前缀标记
    - 重要信息重复编码强化记忆
    - 时序信息编码在特殊token中
    """
    
    # 记忆类型前缀
    TYPE_PREFIX = {
        "chat": "",           # 普通对话
        "fact": "事实:",      # 事实性知识
        "preference": "偏好:", # 用户偏好
        "habit": "习惯:",     # 使用习惯
        "task": "任务:",      # 任务相关
        "emotion": "情感:",   # 情感状态
    }
    
    @staticmethod
    def encode_interaction(
        user_input: str,
        ai_response: str,
        memory_type: str = "chat",
        context: str = "",
        importance: float = 0.5,
    ) -> str:
        """
        编码单次交互为训练文本
        
        关键设计：将记忆编码为明确的问答对格式
        模型学习 "问:xxx 答:xxx" 的模式
        多次重复编码强化记忆（类似人类反复记忆）
        """
        prefix = MemoryEncoder.TYPE_PREFIX.get(memory_type, "")
        
        # 主要格式：问答对
        if context:
            encoded = f"[BOS]{prefix}{context}[SEP]问:{user_input}[SEP]答:{ai_response}[EOS]"
        else:
            encoded = f"[BOS]{prefix}问:{user_input}[SEP]答:{ai_response}[EOS]"
        
        return encoded
    
    @staticmethod
    def encode_fact(fact: str, source: str = "") -> str:
        """编码事实性知识"""
        if source:
            return f"[BOS]事实:来自{source}[SEP]问:{fact}[SEP]答:{fact}[EOS]"
        return f"[BOS]事实:[SEP]问:{fact}[SEP]答:{fact}[EOS]"
    
    @staticmethod
    def encode_preference(category: str, preference: str) -> str:
        """编码用户偏好"""
        return f"[BOS]偏好:{category}[SEP]问:你的{category}是什么[SEP]答:{preference}[EOS]"
    
    @staticmethod
    def encode_habit(trigger: str, behavior: str) -> str:
        """编码使用习惯"""
        return f"[BOS]习惯:[SEP]问:当{trigger}[SEP]答:{behavior}[EOS]"
    
    @staticmethod
    def encode_for_retrieval(query: str, memory_type: str = "") -> str:
        """编码检索查询"""
        prefix = MemoryEncoder.TYPE_PREFIX.get(memory_type, "") if memory_type else ""
        if prefix:
            return f"[BOS]{prefix}问:{query}[SEP]答:"
        return f"[BOS]问:{query}[SEP]答:"


# ============================================================
# 参数化记忆模型核心
# ============================================================

class ParametricMemoryModel:
    """
    参数化记忆模型 — 模型即数据库
    
    工作流程：
    1. 基座模型(SwarmModel) + LoRA适配器
    2. 用户交互 → 记忆编码器 → 训练样本
    3. 积累N条 → 微调LoRA → 记忆写入参数
    4. 检索时直接推理，从参数生成个性化回答
    
    关键优势：
    - 个性化自然流露，不是检索拼接
    - 参数压缩，比存原始文本省空间
    - 长期积累的参数高度贴合用户
    - 推理速度快，一次前向传播
    """
    
    def __init__(
        self,
        model_config: str = "small",     # 基座模型配置
        lora_rank: int = 8,               # LoRA秩
        lora_alpha: float = 16.0,         # LoRA缩放
        max_len: int = 512,               # 最大序列长度（必须匹配预训练权重）
        accumulate_steps: int = 10,       # 积累N条交互后微调
        learning_rate: float = 1e-4,      # 微调学习率
        micro_epochs: int = 3,            # 每次微调训练轮数
        write_mode: str = "instant",      # 'instant'即时写入 / 'batch'批量写入
        device: str = "auto",
        save_dir: str = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "training", "memory_model"),
    ):
        self.model_config = model_config
        self.lora_rank = lora_rank
        self.lora_alpha = lora_alpha
        self.max_len = max_len
        self.accumulate_steps = accumulate_steps
        self.learning_rate = learning_rate
        self.micro_epochs = micro_epochs
        self.write_mode = write_mode
        self.save_dir = save_dir
        
        # 设备
        if device == "auto":
            self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        else:
            self.device = torch.device(device)
        
        # 初始化模型
        self._init_model()
        
        # 记忆缓冲区
        self.memory_buffer = deque(maxlen=1000)  # 待训练的记忆
        self.batch_fallback = deque()  # 即时写入失败时的批量回退队列
        self.total_memories = 0                    # 已写入的记忆总数
        self.total_train_steps = 0                 # 总训练步数
        
        # 训练线程锁
        self._train_lock = threading.Lock()
        self._is_training = False
        
        # 记忆快照历史（用于回滚和版本管理）
        self.snapshot_history = deque(maxlen=10)
        
        # 统计
        self.stats = {
            "memories_stored": 0,
            "memories_written_to_params": 0,
            "train_sessions": 0,
            "total_train_time_sec": 0,
        }
    
    def _init_model(self):
        """初始化基座模型 + LoRA"""
        sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
        from training.model import SwarmModel
        from training.tokenizer import SwarmTokenizer
        
        # 先初始化分词器（确定vocab_size）
        self.tokenizer = SwarmTokenizer(vocab_size=8192)
        self._init_tokenizer()
        
        # 用分词器的实际vocab_size创建模型
        actual_vocab = getattr(self.tokenizer, 'vocab_size_actual', 8192)
        if actual_vocab < 100:
            actual_vocab = 8192  # fallback
        
        # 创建基座模型
        self.base_model = SwarmModel.from_config(
            self.model_config,
            vocab_size=actual_vocab,
            max_len=self.max_len,
        )
        
        # 加载预训练权重（关键：记忆模型需要已预训练的基座）
        self._load_pretrained_weights(actual_vocab)
        
        # 应用LoRA
        self.model, self.lora_layers = apply_lora_to_model(
            self.base_model,
            rank=self.lora_rank,
            alpha=self.lora_alpha,
        )
        self.model = self.model.to(self.device)
        
        # 优化器（只优化LoRA参数）
        lora_params = get_lora_params(self.model)
        self.optimizer = torch.optim.AdamW(lora_params, lr=self.learning_rate, weight_decay=0.01)
        
        # 模型信息
        total_params = sum(p.numel() for p in self.model.parameters())
        trainable_params = sum(p.numel() for p in self.model.parameters() if p.requires_grad)
        
        self.model_info = {
            "total_params": total_params,
            "trainable_params": trainable_params,
            "trainable_ratio": f"{trainable_params/total_params*100:.1f}%",
            "lora_rank": self.lora_rank,
            "lora_layers": len(self.lora_layers),
            "device": str(self.device),
            "vocab_size": actual_vocab,
        }
        
        print(f"[ParametricMemory] 初始化完成")
        print(f"  总参数: {total_params:,}, 可训练: {trainable_params:,} ({trainable_params/total_params*100:.1f}%)")
        print(f"  LoRA层: {len(self.lora_layers)}, 秩: {self.lora_rank}, 词表: {actual_vocab}")
        print(f"  设备: {self.device}")
    
    def _load_pretrained_weights(self, actual_vocab: int):
        """加载预训练权重作为基座（记忆写入的前提）"""
        import glob
        
        # 按优先级查找预训练权重
        weight_paths = [
            os.path.join(_MODEL_DIR, self.model_config, "model.pt"),
            os.path.join(_MODEL_DIR, self.model_config, "model_gpu.pt"),
        ]
        
        loaded = False
        for wpath in weight_paths:
            if os.path.exists(wpath):
                try:
                    checkpoint = torch.load(wpath, map_location=self.device, weights_only=False)
                    
                    if isinstance(checkpoint, dict) and "model_state_dict" in checkpoint:
                        state_dict = checkpoint["model_state_dict"]
                        saved_vocab = checkpoint.get("vocab_size", actual_vocab)
                    elif isinstance(checkpoint, dict):
                        state_dict = checkpoint
                        saved_vocab = actual_vocab
                    else:
                        continue
                    
                    # vocab_size可能不匹配，需要处理
                    if saved_vocab != actual_vocab:
                        print(f"[ParametricMemory] 词表不匹配(存{saved_vocab}/需{actual_vocab})，适配权重...")
                        state_dict = self._adapt_vocab(state_dict, saved_vocab, actual_vocab)
                    
                    # 加载权重
                    missing, unexpected = self.base_model.load_state_dict(state_dict, strict=False)
                    if missing:
                        print(f"[ParametricMemory] 缺失键: {len(missing)}")
                    if unexpected:
                        print(f"[ParametricMemory] 多余键: {len(unexpected)}")
                    
                    loss = checkpoint.get("loss", "?") if isinstance(checkpoint, dict) else "?"
                    print(f"[ParametricMemory] 加载预训练权重: {self.model_config}, loss={loss}")
                    loaded = True
                    break
                    
                except Exception as e:
                    print(f"[ParametricMemory] 加载权重失败({wpath}): {e}")
        
        if not loaded:
            print(f"[ParametricMemory] 未找到预训练权重，使用随机初始化（记忆效果会差）")
    
    def _adapt_vocab(self, state_dict: dict, old_vocab: int, new_vocab: int) -> dict:
        """适配词表大小差异（截断或扩展embedding/lm_head）"""
        adapted = {}
        for key, value in state_dict.items():
            if 'tok_emb.weight' in key or 'head.weight' in key:
                if value.shape[0] != new_vocab:
                    if new_vocab < value.shape[0]:
                        # 截断
                        adapted[key] = value[:new_vocab]
                    else:
                        # 扩展（用随机初始化填充）
                        old_size = value.shape[0]
                        dim = value.shape[1]
                        new_weight = torch.zeros(new_vocab, dim)
                        new_weight[:old_size] = value
                        nn.init.normal_(new_weight[old_size:], mean=0.0, std=0.02)
                        adapted[key] = new_weight
                else:
                    adapted[key] = value
            else:
                adapted[key] = value
        return adapted
    
    def _init_tokenizer(self):
        """初始化分词器：尝试加载已有，否则快速训练"""
        from training.tokenizer import SwarmTokenizer as _SwarmTokenizer
        
        # 尝试加载已有的分词器
        for tok_path in [
            os.path.join(_MODEL_DIR, "small", "tokenizer.json"),
            os.path.join(_MODEL_DIR, "tiny", "tokenizer.json"),
        ]:
            if os.path.exists(tok_path):
                try:
                    loaded = _SwarmTokenizer.load(tok_path)
                    if loaded.vocab_size_actual > 100:
                        self.tokenizer = loaded  # 替换为新加载的实例
                        print(f"[ParametricMemory] 加载已有分词器: {self.tokenizer.vocab_size_actual} tokens")
                        return
                except Exception as e:
                    print(f"[ParametricMemory] 加载分词器失败: {e}")
        
        # 快速训练一个基础分词器
        # 包含常用中文词和特殊标记
        base_texts = []
        # 基础中文词
        base_texts.extend(["你好", "我是虫群助手", "用户", "助手", "事实", "偏好", "习惯", "任务",
                           "情感", "回答", "问题", "帮助", "谢谢", "再见", "编程", "开发", "测试",
                           "部署", "模型", "训练", "数据", "系统", "功能", "项目", "工作", "学习"])
        # 常见对话
        base_texts.extend(["今天天气怎么样", "帮我写一个函数", "什么是虫群", "我喜欢用Python编程",
                           "好的我记住了", "请告诉我更多", "这是一个好主意", "我需要帮助"])
        # 单字覆盖（确保中文基本字都在词表中）
        import string
        common_chars = "的一是不了人我在有他这为之大来以个中上们到说时地也子就道要和去你能对下看行吗着很自会将那给又与从被但让把比等其已或及最更而些只如它为然做方因当所前此两想问此知只使点些因当正新样样心意把比情理相法然体合通量力长电手区计质群位品展复证化任件单据志录养存查调参设层系各部度程表性命定实内三使加系外样问间工式"
        for c in common_chars:
            base_texts.append(c)
        
        self.tokenizer.train(base_texts * 100, min_freq=1)
        print(f"[ParametricMemory] 新建分词器: {self.tokenizer.vocab_size_actual} tokens")
    
    # ============================================================
    # 核心: 存储记忆
    # ============================================================
    
    def store(
        self,
        user_input: str,
        ai_response: str,
        memory_type: str = "chat",
        context: str = "",
        importance: float = 0.5,
    ) -> str:
        """
        存储一条记忆
        
        记忆先进入缓冲区，积累到阈值后自动触发微调写入参数
        重要度高的记忆会被重复编码（强化记忆）
        """
        memory_id = hashlib.md5(
            f"{user_input}{ai_response}{time.time()}".encode()
        ).hexdigest()[:12]
        
        # 编码为训练文本
        encoded_text = MemoryEncoder.encode_interaction(
            user_input, ai_response, memory_type, context, importance
        )
        
        # 记忆条目
        memory_entry = {
            "id": memory_id,
            "type": memory_type,
            "text": encoded_text,
            "importance": importance,
            "timestamp": datetime.now().isoformat(),
            "user_input": user_input[:100],
            "ai_response": ai_response[:100],
        }
        
        # 加入缓冲区
        self.memory_buffer.append(memory_entry)
        self.total_memories += 1
        self.stats["memories_stored"] += 1
        
        # 高重要度记忆：重复编码强化（类比人类反复记忆）
        if importance >= 0.8:
            for _ in range(3):  # 重要记忆额外3次
                self.memory_buffer.append(memory_entry)
        elif importance >= 0.5:
            for _ in range(1):  # 普通记忆额外1次
                self.memory_buffer.append(memory_entry)
        
        # 检查是否需要触发微调
        if len(self.memory_buffer) >= self.accumulate_steps:
            self._trigger_write()
        
        return memory_id
    
    def store_fact(self, fact: str, source: str = "") -> str:
        """存储事实性知识"""
        encoded = MemoryEncoder.encode_fact(fact, source)
        memory_id = hashlib.md5(f"{fact}{time.time()}".encode()).hexdigest()[:12]
        
        self.memory_buffer.append({
            "id": memory_id,
            "type": "fact",
            "text": encoded,
            "importance": 0.7,  # 事实通常重要
            "timestamp": datetime.now().isoformat(),
        })
        self.total_memories += 1
        self.stats["memories_stored"] += 1
        
        if len(self.memory_buffer) >= self.accumulate_steps:
            self._trigger_write()
        
        return memory_id
    
    def store_preference(self, category: str, preference: str) -> str:
        """存储用户偏好"""
        encoded = MemoryEncoder.encode_preference(category, preference)
        memory_id = hashlib.md5(f"{category}{preference}{time.time()}".encode()).hexdigest()[:12]
        
        self.memory_buffer.append({
            "id": memory_id,
            "type": "preference",
            "text": encoded,
            "importance": 0.9,  # 偏好非常重要
            "timestamp": datetime.now().isoformat(),
        })
        self.total_memories += 1
        self.stats["memories_stored"] += 1
        
        if len(self.memory_buffer) >= self.accumulate_steps:
            self._trigger_write()
        
        return memory_id
    
    # ============================================================
    # 核心: 记忆写入参数（微调）
    # ============================================================
    
    # ============================================================
    # 核心写入方法：即时写入 vs 批量写入
    # ============================================================
    
    def _instant_write(self, memory_entry: Dict):
        """
        即时记忆写入 (Instant Memory Write)
        
        核心理念：一次交互，一步更新，立刻记住
        区别于传统微调（积累一批→多轮训练），即时写入：
        - 每条记忆只做1步梯度更新
        - 学习率按记忆重要度动态缩放
        - 写入后立即验证，不通过则追加1步
        - 像人脑一样，经历一次就形成记忆痕迹
        
        适用场景：手机/边缘设备24小时后台运行
        - 用户每次对话后触发一次即时写入
        - 不阻塞交互，后台线程执行
        - 长期积累后模型自然贴合个人习惯
        """
        tokens = self.tokenizer.encode(memory_entry["text"], add_special=False)
        if len(tokens) < 4:
            return False
        
        # 找回答起始位置
        sep_tokens = self.tokenizer.encode("[SEP]", add_special=False)
        answer_start = 0
        for i in range(len(tokens) - len(sep_tokens) + 1):
            if tokens[i:i+len(sep_tokens)] == sep_tokens:
                answer_start = i + len(sep_tokens)
        
        if len(tokens) > self.max_len - 1:
            tokens = tokens[:self.max_len - 1]
        
        input_ids = tokens
        targets = tokens[1:] + [0]
        for i in range(min(answer_start, len(targets))):
            targets[i] = -100
        
        pad_id = self.tokenizer.pad_id
        pad_len = self.max_len - len(input_ids)
        input_ids = input_ids + [pad_id] * pad_len
        targets = targets + [-100] * pad_len
        input_ids = input_ids[:self.max_len]
        targets = targets[:self.max_len]
        
        # 转tensor
        input_t = torch.tensor([input_ids], dtype=torch.long, device=self.device)
        target_t = torch.tensor([targets], dtype=torch.long, device=self.device)
        
        # 动态学习率：重要度越高，步长越大
        importance = memory_entry.get("importance", 0.5)
        instant_lr = self.learning_rate * (1.0 + importance * 3.0)  # 重要记忆3倍学习率
        max_steps = 3 if importance >= 0.8 else 2  # 重要记忆最多3步
        
        self.model.train()
        written = False
        
        for step in range(max_steps):
            logits, loss = self.model(input_t, targets=target_t)
            
            # 对比损失：防止模型学到通用回答
            # 如果loss已经很低，说明模型已经记住了
            if loss.item() < 1.0:
                written = True
                break
            
            self.optimizer.zero_grad()
            loss.backward()
            torch.nn.utils.clip_grad_norm_(get_lora_params(self.model), 1.0)
            
            # 即时写入用更大的步长
            for pg in self.optimizer.param_groups:
                old_lr = pg['lr']
                pg['lr'] = instant_lr
            self.optimizer.step()
            for pg in self.optimizer.param_groups:
                pg['lr'] = old_lr
            
            self.total_train_steps += 1
            
            # 验证：生成看看是否记住了
            if step < max_steps - 1:  # 最后一步不需要验证
                self.model.eval()
                verify_result = self._verify_memory(memory_entry)
                self.model.train()
                if verify_result:
                    written = True
                    break
        
        return written
    
    def _verify_memory(self, memory_entry: Dict) -> bool:
        """
        验证记忆是否已成功写入参数
        
        方法：用查询生成回答，检查是否包含期望的关键词
        这是一种轻量级验证，不要求完全匹配
        """
        text = memory_entry["text"]
        # 从编码文本中提取问答
        if "[SEP]答:" in text:
            parts = text.split("[SEP]答:")
            if len(parts) >= 2:
                expected_answer = parts[-1].replace("[EOS]", "").strip()
                # 提取关键词（取前几个有意义的词）
                keywords = [w for w in expected_answer if len(w) >= 2][:3]
                
                if not keywords:
                    return False
                
                # 从问的部分提取查询
                query_part = parts[0].split("[SEP]问:")[-1] if "[SEP]问:" in parts[0] else ""
                if not query_part:
                    return False
                
                # 生成回答
                query_encoded = MemoryEncoder.encode_for_retrieval(query_part)
                query_tokens = self.tokenizer.encode(query_encoded, add_special=False)
                if len(query_tokens) > self.max_len - 32:
                    query_tokens = query_tokens[:self.max_len - 32]
                
                input_ids = torch.tensor([query_tokens], dtype=torch.long, device=self.device)
                with torch.no_grad():
                    output = self.model.generate(input_ids, max_new_tokens=32, 
                                                 temperature=0.3, top_k=10,
                                                 eos_id=self.tokenizer.eos_id)
                new_tokens = output[0].tolist()[len(query_tokens):]
                generated = self.tokenizer.decode(new_tokens)
                
                # 检查是否包含关键词
                match_count = sum(1 for kw in keywords if kw in generated)
                return match_count >= len(keywords) * 0.5  # 至少匹配一半关键词
        
        return False
    
    # ============================================================
    # 写入调度：即时模式 vs 批量模式
    # ============================================================
    
    def _trigger_write(self):
        """
        触发记忆写入
        
        两种模式：
        1. 即时模式(write_mode='instant')：每条记忆立即写入，适合手机/边缘设备
        2. 批量模式(write_mode='batch')：积累后批量训练，适合服务器/有GPU时
        """
        if self._is_training:
            return
        
        if self.write_mode == 'instant':
            # 即时模式：在后台逐条写入
            thread = threading.Thread(target=self._instant_write_loop, daemon=True)
            thread.start()
        else:
            # 批量模式：在后台批量训练
            thread = threading.Thread(target=self._write_memories_to_params, daemon=True)
            thread.start()
    
    def _instant_write_loop(self):
        """即时写入循环：逐条处理缓冲区中的记忆"""
        with self._train_lock:
            if self._is_training or len(self.memory_buffer) == 0:
                return
            
            self._is_training = True
            start_time = time.time()
            written_count = 0
            
            try:
                # 保存快照
                snapshot = get_lora_state(self.model)
                self.snapshot_history.append({
                    "state": snapshot,
                    "timestamp": datetime.now().isoformat(),
                    "memories": self.total_memories,
                })
                
                while self.memory_buffer:
                    mem = self.memory_buffer.popleft()
                    success = self._instant_write(mem)
                    if success:
                        written_count += 1
                        self.stats["instant_writes"] = self.stats.get("instant_writes", 0) + 1
                    else:
                        # 即时写入失败，放回批量训练队列
                        self.batch_fallback.append(mem)
                
                # 批量回退的记忆用传统方式训练
                if self.batch_fallback:
                    for mem in self.batch_fallback:
                        self.memory_buffer.append(mem)
                    self.batch_fallback.clear()
                    if self.memory_buffer:
                        self._write_memories_to_params_internal()
                
                elapsed = time.time() - start_time
                self.stats["memories_written_to_params"] += written_count
                self.stats["train_sessions"] += 1
                self.stats["total_train_time_sec"] += elapsed
                
                print(f"[ParametricMemory] 即时写入: {written_count}条记忆, "
                      f"回退{len(self.batch_fallback)}条, 用时{elapsed:.1f}s")
                
            except Exception as e:
                print(f"[ParametricMemory] 即时写入失败: {e}")
                if self.snapshot_history:
                    last_snapshot = self.snapshot_history[-1]
                    set_lora_state(self.model, last_snapshot["state"])
            
            finally:
                self._is_training = False
    
    def _write_memories_to_params_internal(self):
        """批量写入的内部实现（不触发线程）"""
        batch_memories = list(self.memory_buffer)
        self.memory_buffer.clear()
        
        train_samples = []
        for mem in batch_memories:
            tokens = self.tokenizer.encode(mem["text"], add_special=False)
            if len(tokens) < 4:
                continue
            
            sep_tokens = self.tokenizer.encode("[SEP]", add_special=False)
            answer_start = 0
            for i in range(len(tokens) - len(sep_tokens) + 1):
                if tokens[i:i+len(sep_tokens)] == sep_tokens:
                    answer_start = i + len(sep_tokens)
            
            if len(tokens) > self.max_len - 1:
                tokens = tokens[:self.max_len - 1]
            
            input_ids = tokens
            targets = tokens[1:] + [0]
            for i in range(min(answer_start, len(targets))):
                targets[i] = -100
            
            pad_id = self.tokenizer.pad_id
            pad_len = self.max_len - len(input_ids)
            input_ids = input_ids + [pad_id] * pad_len
            targets = targets + [-100] * pad_len
            input_ids = input_ids[:self.max_len]
            targets = targets[:self.max_len]
            
            train_samples.append({"input_ids": input_ids, "targets": targets})
        
        if not train_samples:
            return
        
        self.model.train()
        best_loss = float('inf')
        for epoch in range(self.micro_epochs):
            epoch_loss = 0.0
            num_batches = 0
            batch_size = min(8, len(train_samples))
            for i in range(0, len(train_samples), batch_size):
                batch = train_samples[i:i+batch_size]
                input_ids = torch.tensor([s["input_ids"] for s in batch], dtype=torch.long, device=self.device)
                targets = torch.tensor([s["targets"] for s in batch], dtype=torch.long, device=self.device)
                logits, loss = self.model(input_ids, targets=targets)
                self.optimizer.zero_grad()
                loss.backward()
                torch.nn.utils.clip_grad_norm_(get_lora_params(self.model), 1.0)
                self.optimizer.step()
                epoch_loss += loss.item()
                num_batches += 1
                self.total_train_steps += 1
            avg_loss = epoch_loss / max(num_batches, 1)
            if avg_loss < best_loss:
                best_loss = avg_loss
        
        self.stats["memories_written_to_params"] += len(batch_memories)
        print(f"[ParametricMemory] 批量写入: {len(batch_memories)}条, loss={best_loss:.4f}")
    
    def _write_memories_to_params(self):
        """
        将缓冲区的记忆微调写入LoRA参数
        
        这是"模型即数据库"的核心实现：
        - 每条记忆是一个训练样本
        - 微调LoRA参数 = 将记忆编码进模型
        - 之后推理时，模型自然"记住"了这些内容
        """
        with self._train_lock:
            if self._is_training or len(self.memory_buffer) == 0:
                return
            
            self._is_training = True
            start_time = time.time()
            
            try:
                # 1. 保存当前快照（用于回滚）
                snapshot = get_lora_state(self.model)
                self.snapshot_history.append({
                    "state": snapshot,
                    "timestamp": datetime.now().isoformat(),
                    "memories": self.total_memories,
                })
                
                # 2. 准备训练数据
                train_samples = []
                batch_memories = list(self.memory_buffer)
                self.memory_buffer.clear()
                
                for mem in batch_memories:
                    # 编码为token
                    tokens = self.tokenizer.encode(mem["text"], add_special=False)
                    if len(tokens) < 4:
                        continue  # 太短跳过
                    
                    # 找到"答:"的位置——只对回答部分计算loss
                    # 编码"答:"的token序列
                    answer_prefix_tokens = self.tokenizer.encode("答:", add_special=False)
                    sep_tokens = self.tokenizer.encode("[SEP]", add_special=False)
                    
                    # 在tokens中找到最后一个[SEP]之后的位置作为回答开始
                    answer_start = 0
                    for i in range(len(tokens) - len(sep_tokens) + 1):
                        if tokens[i:i+len(sep_tokens)] == sep_tokens:
                            answer_start = i + len(sep_tokens)
                    
                    # 截断
                    if len(tokens) > self.max_len - 1:
                        tokens = tokens[:self.max_len - 1]
                    
                    # 输入和目标
                    input_ids = tokens
                    targets = tokens[1:] + [0]
                    
                    # 对回答之前的部分，target设为-100（不计算loss）
                    # 这样模型只学习回答部分的预测
                    for i in range(min(answer_start, len(targets))):
                        targets[i] = -100
                    
                    # padding
                    pad_id = self.tokenizer.pad_id
                    pad_len = self.max_len - len(input_ids)
                    input_ids = input_ids + [pad_id] * pad_len
                    targets = targets + [-100] * pad_len
                    
                    # 截断
                    input_ids = input_ids[:self.max_len]
                    targets = targets[:self.max_len]
                    
                    train_samples.append({
                        "input_ids": input_ids,
                        "targets": targets,
                    })
                
                if not train_samples:
                    self._is_training = False
                    return
                
                # 3. 微调训练
                self.model.train()
                best_loss = float('inf')
                
                for epoch in range(self.micro_epochs):
                    epoch_loss = 0.0
                    num_batches = 0
                    
                    # 简单batch训练
                    batch_size = min(8, len(train_samples))
                    for i in range(0, len(train_samples), batch_size):
                        batch = train_samples[i:i+batch_size]
                        
                        input_ids = torch.tensor(
                            [s["input_ids"] for s in batch], 
                            dtype=torch.long, device=self.device
                        )
                        targets = torch.tensor(
                            [s["targets"] for s in batch], 
                            dtype=torch.long, device=self.device
                        )
                        
                        logits, loss = self.model(input_ids, targets=targets)
                        
                        self.optimizer.zero_grad()
                        loss.backward()
                        torch.nn.utils.clip_grad_norm_(get_lora_params(self.model), 1.0)
                        self.optimizer.step()
                        
                        epoch_loss += loss.item()
                        num_batches += 1
                        self.total_train_steps += 1
                    
                    avg_loss = epoch_loss / max(num_batches, 1)
                    
                    if avg_loss < best_loss:
                        best_loss = avg_loss
                
                # 4. 训练完成
                elapsed = time.time() - start_time
                self.stats["memories_written_to_params"] += len(batch_memories)
                self.stats["train_sessions"] += 1
                self.stats["total_train_time_sec"] += elapsed
                
                print(f"[ParametricMemory] 写入完成: {len(batch_memories)}条记忆, "
                      f"loss={best_loss:.4f}, 用时{elapsed:.1f}s")
                
            except Exception as e:
                print(f"[ParametricMemory] 写入失败: {e}")
                # 回滚到上一个快照
                if self.snapshot_history:
                    last_snapshot = self.snapshot_history[-1]
                    set_lora_state(self.model, last_snapshot["state"])
                    print(f"[ParametricMemory] 已回滚到快照")
            
            finally:
                self._is_training = False
    
    # ============================================================
    # 核心: 记忆检索（从参数生成）
    # ============================================================
    
    def recall(self, query: str, memory_type: str = "", max_tokens: int = 64,
               temperature: float = 0.7) -> Dict:
        """
        记忆检索 — 从模型参数生成回答
        
        与数据库检索的本质区别：
        - 数据库：匹配相似文本 → 返回原文片段
        - 参数化：模型基于学到的参数 → 直接生成个性化回答
        """
        self.model.eval()
        
        # 编码查询
        query_encoded = MemoryEncoder.encode_for_retrieval(query, memory_type)
        tokens = self.tokenizer.encode(query_encoded, add_special=False)
        
        if len(tokens) == 0:
            return {"response": "", "confidence": 0.0, "source": "parametric"}
        
        # 截断
        if len(tokens) > self.max_len - max_tokens:
            tokens = tokens[:self.max_len - max_tokens]
        
        input_ids = torch.tensor([tokens], dtype=torch.long, device=self.device)
        
        # 生成
        with torch.no_grad():
            output_ids = self.model.generate(
                input_ids,
                max_new_tokens=max_tokens,
                temperature=temperature,
                top_k=40,
                eos_id=self.tokenizer.eos_id,
            )
        
        # 解码（只取新生成的部分）
        new_tokens = output_ids[0].tolist()[len(tokens):]
        response = self.tokenizer.decode(new_tokens)
        
        # 清理
        response = self._clean_response(response)
        
        return {
            "response": response,
            "confidence": self._estimate_confidence(response),
            "source": "parametric",
            "query": query,
            "model_info": self.model_info,
        }
    
    def recall_with_context(self, query: str, context: str = "", 
                           max_tokens: int = 64) -> Dict:
        """带上下文的记忆检索"""
        full_query = f"{context} {query}" if context else query
        return self.recall(full_query, max_tokens=max_tokens)
    
    def _clean_response(self, text: str) -> str:
        """清理生成文本"""
        # 去特殊token
        for tok in ["[PAD]", "[UNK]", "[BOS]", "[EOS]", "[CLS]", "[SEP]", "</w>"]:
            text = text.replace(tok, "")
        # 去多余空格
        text = re.sub(r'\s+', ' ', text).strip()
        return text
    
    def _estimate_confidence(self, response: str) -> float:
        """简单评估回答置信度"""
        if not response:
            return 0.0
        # 基于长度和完整度
        length_score = min(len(response) / 20.0, 1.0)
        # 有明确回答的标记
        has_answer = any(kw in response for kw in ["是", "的", "了", "可以", "因为"])
        confidence = length_score * 0.5 + (0.5 if has_answer else 0.0)
        return min(confidence, 1.0)
    
    # ============================================================
    # 持久化
    # ============================================================
    
    def save(self, path: str = None):
        """保存参数化记忆模型"""
        path = path or self.save_dir
        os.makedirs(path, exist_ok=True)
        
        # 保存LoRA参数
        lora_state = get_lora_state(self.model)
        torch.save(lora_state, os.path.join(path, "lora_params.pt"))
        
        # 保存基座模型
        base_state = {k: v for k, v in self.model.state_dict().items() 
                      if 'lora' not in k}
        torch.save(base_state, os.path.join(path, "base_model.pt"))
        
        # 保存分词器
        self.tokenizer.save(os.path.join(path, "tokenizer.json"))
        
        # 保存元信息
        meta = {
            "model_config": self.model_config,
            "lora_rank": self.lora_rank,
            "lora_alpha": self.lora_alpha,
            "max_len": self.max_len,
            "total_memories": self.total_memories,
            "total_train_steps": self.total_train_steps,
            "stats": self.stats,
            "model_info": self.model_info,
            "saved_at": datetime.now().isoformat(),
        }
        with open(os.path.join(path, "meta.json"), "w") as f:
            json.dump(meta, f, ensure_ascii=False, indent=2)
        
        print(f"[ParametricMemory] 已保存到 {path}")
    
    def load(self, path: str = None):
        """加载参数化记忆模型"""
        path = path or self.save_dir
        
        # 加载LoRA参数
        lora_path = os.path.join(path, "lora_params.pt")
        if os.path.exists(lora_path):
            lora_state = torch.load(lora_path, map_location=self.device)
            set_lora_state(self.model, lora_state)
        
        # 加载分词器
        tok_path = os.path.join(path, "tokenizer.json")
        if os.path.exists(tok_path):
            self.tokenizer.load(tok_path)
        
        # 加载元信息
        meta_path = os.path.join(path, "meta.json")
        if os.path.exists(meta_path):
            with open(meta_path) as f:
                meta = json.load(f)
            self.total_memories = meta.get("total_memories", 0)
            self.total_train_steps = meta.get("total_train_steps", 0)
            self.stats = meta.get("stats", self.stats)
        
        print(f"[ParametricMemory] 已加载: {self.total_memories}条记忆, {self.total_train_steps}步训练")
    
    # ============================================================
    # 信息与调试
    # ============================================================
    
    def get_status(self) -> Dict:
        """获取记忆模型状态"""
        return {
            "model_info": self.model_info,
            "buffer_size": len(self.memory_buffer),
            "total_memories": self.total_memories,
            "total_train_steps": self.total_train_steps,
            "is_training": self._is_training,
            "stats": self.stats,
            "snapshot_count": len(self.snapshot_history),
        }
    
    def force_write(self):
        """强制写入所有缓冲区记忆（不管阈值）"""
        if len(self.memory_buffer) > 0:
            self._write_memories_to_params()
    
    def rollback(self):
        """回滚到上一个快照"""
        if self.snapshot_history:
            last = self.snapshot_history.pop()
            set_lora_state(self.model, last["state"])
            print(f"[ParametricMemory] 已回滚到 {last['timestamp']}")
        else:
            print("[ParametricMemory] 无快照可回滚")


# ============================================================
# 记忆蒸馏 — 防止参数膨胀
# ============================================================

class MemoryDistiller:
    """
    记忆蒸馏器 — 定期将LoRA参数蒸馏回基座
    
    问题：LoRA参数持续增长，可能偏离基座太远
    方案：定期将LoRA参数合并回基座权重，然后重置LoRA
    
    W_new = W_base + (B @ A) * scaling
    然后重置 B=0, A=kaiming
    """
    
    @staticmethod
    def distill(model: nn.Module) -> int:
        """将LoRA参数合并回基座，返回合并的层数"""
        merged_count = 0
        
        for name, module in model.named_modules():
            if isinstance(module, LoRALayer):
                # 合并: W = W + B @ A * scaling
                with torch.no_grad():
                    delta_w = (module.lora_B @ module.lora_A) * module.scaling
                    module.original.weight.data += delta_w
                    
                    # 重置LoRA
                    nn.init.kaiming_uniform_(module.lora_A, a=math.sqrt(5))
                    nn.init.zeros_(module.lora_B)
                
                merged_count += 1
        
        print(f"[MemoryDistiller] 蒸馏完成: 合并{merged_count}层LoRA回基座")
        return merged_count


# ============================================================
# 快速测试
# ============================================================

if __name__ == "__main__":
    print("=" * 60)
    print("虫群v8 — 参数化记忆模型 测试")
    print("=" * 60)
    
    # 1. 创建记忆模型
    print("\n[1] 初始化参数化记忆模型...")
    pm = ParametricMemoryModel(
        model_config="tiny",  # 测试用tiny（小模型快速验证）
        lora_rank=4,
        accumulate_steps=5,   # 5条触发写入
        micro_epochs=5,       # 多训练几轮加深记忆
    )
    
    # 2. 存储记忆
    print("\n[2] 存储记忆...")
    memories = [
        ("你叫什么名字", "我是虫群助手，你的个人AI"),
        ("我喜欢用Python编程", "好的，我记住了你喜欢Python"),
        ("今天天气怎么样", "今天是晴天，适合出门"),
        ("帮我写一个函数", "好的，请告诉我函数的功能需求"),
        ("什么是虫群", "虫群是分布式小模型聚合系统"),
        ("我住在北京", "好的，我记住了你住在北京"),
        ("我的工作是什么", "你是一名软件开发工程师"),
        ("我喜欢的食物", "你喜欢川菜和火锅"),
        ("我的项目叫什么", "你的项目叫虫群Swarm"),
        ("我最近在忙什么", "你最近在开发参数化记忆模型"),
    ]
    
    for user, ai in memories:
        mid = pm.store(user, ai, memory_type="chat")
        print(f"  存储: {user} → {mid}")
    
    # 3. 等待写入完成
    print("\n[3] 等待记忆写入...")
    time.sleep(3)
    
    # 4. 存储偏好
    print("\n[4] 存储偏好...")
    pm.store_preference("编程语言", "Python")
    pm.store_preference("工作风格", "简洁高效")
    pm.store_fact("虫群项目始于2026年5月", source="项目记录")
    
    # 强制写入
    pm.force_write()
    
    # 5. 检索记忆
    print("\n[5] 检索记忆...")
    queries = ["我喜欢什么", "虫群是什么", "我叫什么"]
    for q in queries:
        result = pm.recall(q, max_tokens=32)
        print(f"  Q: {q}")
        print(f"  A: {result['response'][:50]} (置信度: {result['confidence']:.2f})")
    
    # 6. 状态
    print("\n[6] 记忆模型状态:")
    status = pm.get_status()
    for k, v in status.items():
        print(f"  {k}: {v}")
    
    print("\n" + "=" * 60)
    print("测试完成!")