#!/usr/bin/env python3 """ 虫群Brain v2.0 — AreaNetwork适配层 包装Meta Model的AreaNetwork(6区×3柱=18柱)，补充chat/stats/get_area_weights等接口 让node_server和chat_engine无需修改即可使用完整版Meta Model引擎 """ import numpy as np import base64 import io import os from typing import Dict, Optional from .functional_area import AreaNetwork, AreaType from .language_decoder import LanguageDecoder class Brain: """ 虫群Brain v2.0 — AreaNetwork适配器 支持两种推理模式: - 原始模式: 全链路forward(6区随机权重, 信号衰减严重) - 映射模式: W_s(sensory直接映射) + P(motor投影), 跳过随机中间层 映射模式在训练后自动启用 """ def __init__(self, dim: int = 75, config: dict = None): self.dim = dim self._config = config # 保存配置供后续使用 # AreaNetwork需要 {'sensory': {...}, ...} 格式 area_config = config.get('areas', config) if config else None self._net = AreaNetwork(config=area_config) self._forward_count = 0 # 默认冻结: forward只做推理，learn/train时临时解冻 self._net.freeze() # 训练映射矩阵(加载后启用映射模式) self._W_s = None # (75, 300) sensory直接映射 self._P = None # (300, 75) motor投影 self._load_mappings() # QA检索索引(训练数据的问题向量→答案) self._qa_vecs = None # (N, 75) 问题编码向量 self._qa_answers = [] # 对应答案列表 self._load_qa_index() # 语言解码层 — motor输出→自回归token生成 self._decoder = None self._init_decoder() def _init_decoder(self): """初始化语言解码层，加载词表和已训练权重""" try: vocab_path = os.path.join(os.path.dirname(__file__), '..', '..', 'models', 'vocab75_clean_v2.pkl') # 尝试json格式(优先) json_path = vocab_path.replace('.pkl', '.json') chars = None if os.path.exists(json_path): import json as _json with open(json_path, 'r', encoding='utf-8') as f: vocab_data = _json.load(f) chars = list(vocab_data) if isinstance(vocab_data, list) else list(vocab_data.keys()) elif os.path.exists(vocab_path): import pickle with open(vocab_path, 'rb') as f: vocab_data = pickle.load(f) chars = list(vocab_data.keys()) if isinstance(vocab_data, dict) else list(vocab_data) else: # 尝试clean词表 alt_path = vocab_path.replace('vocab75_clean_v2', 'vocab75_clean') alt_json = alt_path.replace('.pkl', '_words.json') if os.path.exists(alt_json): import json as _json with open(alt_json, 'r', encoding='utf-8') as f: chars = _json.load(f) elif os.path.exists(alt_path.replace('.pkl', '.npz')): # 从npz加载但需要词表 pass if chars is None: print('[Brain] 未找到词表，跳过解码层初始化') self._decoder = None return motor_dim = self.config.get('motor_output_dim', 300) hidden_dim = self.config.get('areas', {}).get('motor', {}).get('hidden_dim', 300) self._decoder = LanguageDecoder(motor_dim=motor_dim, hidden_dim=hidden_dim, vocab_size=len(chars) + 3) self._decoder.set_vocab(chars) # 尝试加载已训练权重 dec_path = os.path.join(os.path.dirname(__file__), '..', '..', 'models', 'decoder_weights.npz') if os.path.exists(dec_path): self._decoder.load(dec_path) print(f"[Brain] 语言解码层已加载训练权重") else: print(f"[Brain] 语言解码层已初始化(vocab={len(chars)}字, 未训练)") except Exception as e: print(f"[Brain] 语言解码层初始化失败: {e}") self._decoder = None @property def areas(self) -> Dict: """兼容旧接口: 返回 {区名: FunctionalArea}""" return {at.value: area for at, area in self._net.areas.items()} def _load_mappings(self): """加载训练映射矩阵(W_s和P), 存在则启用映射模式""" try: import os path = os.path.join(os.path.dirname(__file__), '..', '..', 'models', 'trained_mappings.npz') if os.path.exists(path): data = np.load(path, allow_pickle=True) self._W_s = data['W_sensory'] # (75, 300) self._P = data['P_motor'] # (300, 75) print(f"[Brain] 映射模式已启用: W_s={self._W_s.shape}, P={self._P.shape}") except Exception as e: print(f"[Brain] 映射矩阵未找到, 使用原始模式: {e}") def _load_qa_index(self): """加载QA检索索引: 对训练数据的问题编码，建立向量检索""" import os, json try: qa_path = os.path.join(os.path.dirname(__file__), '..', '..', 'data', 'qa_training.json') if not os.path.exists(qa_path): print(f"[Brain] QA训练数据未找到: {qa_path}") return with open(qa_path, 'r') as f: qa_data = json.load(f) from .semantic_encoder import get_encoder encoder = get_encoder() vecs, answers = [], [] for item in qa_data: q = item.get('q') or item.get('question', '') a = item.get('a') or item.get('answer', '') if not q or not a: continue v = encoder.encode(q) if v is not None and np.linalg.norm(v) > 0: vecs.append(v[:self.dim]) answers.append(a) if vecs: self._qa_vecs = np.array(vecs, dtype=np.float32) self._qa_answers = answers print(f"[Brain] QA索引已加载: {len(answers)}条, vecs={self._qa_vecs.shape}") except Exception as e: print(f"[Brain] QA索引加载失败: {e}") def _qa_retrieve(self, input_vec, threshold=0.75): """用输入向量检索最匹配的QA对""" if self._qa_vecs is None or len(self._qa_answers) == 0: return None, 0.0 v = input_vec[:self.dim] v_norm = np.linalg.norm(v) if v_norm < 1e-8: return None, 0.0 # 批量cosine norms = np.linalg.norm(self._qa_vecs, axis=1, keepdims=True) norms = np.maximum(norms, 1e-8) sims = (self._qa_vecs @ v) / (norms.ravel() * v_norm) best_idx = int(np.argmax(sims)) best_sim = float(sims[best_idx]) if best_sim >= threshold: return self._qa_answers[best_idx], best_sim return None, best_sim def forward(self, input_vec) -> Dict[str, np.ndarray]: """ 前向传播 — 6区协作推理 映射模式: 输入75维→W_s(75×300)→sensory→AreaNetwork全链路→motor(300)→P(300×75)→输出75维 原始模式: 输入75维→AreaNetwork全链路→motor(300) """ if isinstance(input_vec, str): return self.chat(input_vec) x = np.asarray(input_vec, dtype=np.float32).ravel()[:self.dim] if len(x) < self.dim: x = np.pad(x, (0, self.dim - len(x))) # AreaNetwork原始forward output, meta = self._net.forward(x) self._forward_count += 1 # 从meta提取各区输出 result = {} area_outputs = meta.get('area_outputs', {}) for area_type, area_data in area_outputs.items(): key = area_type.value if hasattr(area_type, 'value') else str(area_type) result[key] = area_data # 映射模式: 用W_s覆盖sensory, 用P投影motor if self._W_s is not None: # W_s直接映射: 输入→sensory(跳过随机sensory层) result['sensory'] = (x @ self._W_s).astype(np.float32) # 确保有motor输出 if 'motor' not in result: result['motor'] = output result['_final'] = output return result def chat(self, text: str) -> Optional[Dict]: """ 文字输入 — 语义编码为75维向量后forward Args: text: 输入文字 Returns: {'text': 回复文字, 'confidence': 置信度, 'areas': 区激活信息} """ # 语义编码（vecs75查表，替代hash） from .semantic_encoder import get_encoder encoder = get_encoder() vec = encoder.encode(text) # QA检索: 编码向量匹配训练数据 # cos>=0.75精确匹配, 0.55~0.75模糊匹配(标注低置信度) qa_answer, qa_sim = self._qa_retrieve(vec, threshold=0.55) if qa_answer: mode = 'qa_brain' if qa_sim >= 0.75 else 'qa_brain_fuzzy' return { 'text': qa_answer, 'confidence': float(qa_sim), 'areas': ['qa_retrieve'], 'decoded_words': [(qa_answer[:8], qa_sim)], 'mode': mode, } result = self.forward(vec) # 从多个区域信号综合计算置信度 confidences = [] for key in ['sensory', 'association', 'prefrontal', 'motor']: arr = result.get(key) if isinstance(arr, np.ndarray) and arr.size > 0: confidences.append(float(np.max(np.abs(arr)))) confidence = max(confidences) if confidences else 0.0 # 语言解码层: motor 300维 → 自回归生成句子 motor_vec = result.get('motor', result.get('_final', np.zeros(1))) if self._decoder is not None and motor_vec.size >= 300: generated = self._decoder.decode(motor_vec[:300], top_k=10) if generated and len(generated) > 1: return { 'text': generated, 'confidence': confidence, 'areas': list(result.keys()), 'decoded_words': [(generated[:8], confidence)], 'mode': 'language_decoder', } # 回退: 最近邻词解码(未训练时) if self._W_s is not None: sensory = result.get('sensory', np.zeros(1)) decode_vec = sensory[:75] if sensory.size >= 75 else motor_vec[:min(75, motor_vec.size)] elif self._P is not None and motor_vec.size >= 300: decode_vec = (motor_vec[:300] @ self._P).astype(np.float32) else: decode_vec = motor_vec top_words = encoder.decode_nearest(decode_vec, top_k=5) if top_words and top_words[0][1] > 0.3: decoded = ' '.join(w for w, s in top_words[:3] if s > 0.3) else: decoded = f'置信度:{confidence:.3f}' return { 'text': f'[Brain] {decoded}', 'confidence': confidence, 'areas': list(result.keys()), 'decoded_words': top_words[:5], } def stats(self) -> Dict: """返回统计信息""" area_info = {} total_params = 0 for area_type, area in self._net.areas.items(): name = area_type.value col_count = len(area.columns) mc_count = sum(len(c.micro_columns) for c in area.columns) # 用微柱的total_params属性统计 params = 0 for col in area.columns: for mc in col.micro_columns: if hasattr(mc, 'total_params'): params += mc.total_params total_params += params area_info[name] = { 'area_name': name, 'columns': col_count, 'micro_columns': mc_count, 'params': params, } return { 'areas': len(self._net.areas), 'total_params': total_params, 'area_details': area_info, 'forward_count': self._forward_count, } def get_area_weights(self, area_name: str) -> np.ndarray: """导出指定区权重 — 深度收集所有numpy属性(含_synaptic内部)""" area_type = AreaType(area_name) area = self._net.areas[area_type] weights_list = [] for col in area.columns: for mc in col.micro_columns: # 收集mc自身的所有numpy属性(含下划线) mc._weight_attrs = [] for attr_name in dir(mc): val = getattr(mc, attr_name, None) if isinstance(val, np.ndarray) and val.ndim >= 1 and val.size < 1_000_000: weights_list.append(val.ravel()) mc._weight_attrs.append(attr_name) # 深入_synaptic收集 syn = getattr(mc, '_synaptic', None) if syn is not None: syn._weight_attrs = [] for attr_name in dir(syn): val = getattr(syn, attr_name, None) if isinstance(val, np.ndarray) and val.ndim >= 1 and val.size < 1_000_000: weights_list.append(val.ravel()) syn._weight_attrs.append(attr_name) if weights_list: return np.concatenate(weights_list) return np.array([], dtype=np.float32) def set_area_weights(self, area_name: str, weights: np.ndarray): """导入指定区权重 — 按记录的属性顺序回填(含_synaptic内部)""" area_type = AreaType(area_name) area = self._net.areas[area_type] offset = 0 for col in area.columns: for mc in col.micro_columns: # 回填mc自身属性 for attr_name in getattr(mc, '_weight_attrs', []): val = getattr(mc, attr_name, None) if isinstance(val, np.ndarray) and val.ndim >= 1: n = val.size if offset + n <= weights.size: setattr(mc, attr_name, weights[offset:offset+n].reshape(val.shape)) offset += n # 回填_synaptic属性 syn = getattr(mc, '_synaptic', None) if syn is not None: for attr_name in getattr(syn, '_weight_attrs', []): val = getattr(syn, attr_name, None) if isinstance(val, np.ndarray) and val.ndim >= 1: n = val.size if offset + n <= weights.size: setattr(syn, attr_name, weights[offset:offset+n].reshape(val.shape)) offset += n def train_decoder(self, qa_data: list, epochs: int = 10, lr: float = 0.01) -> Dict: """训练语言解码层 — 用QA数据teacher forcing训练W_out映射 Args: qa_data: [{'q': 问题, 'a': 答案}, ...] epochs: 训练轮数 lr: 学习率 """ if self._decoder is None: return {'error': '解码层未初始化'} from .semantic_encoder import get_encoder encoder = get_encoder() total_loss = 0.0 n_trained = 0 for epoch in range(epochs): epoch_loss = 0.0 for item in qa_data: q = item.get('q') or item.get('question', '') a = item.get('a') or item.get('answer', '') if not q or not a: continue # 问题编码 → forward获取motor输出 vec = encoder.encode(q) if vec is None: continue result = self.forward(vec) motor_vec = result.get('motor', np.zeros(self.config.get('motor_output_dim', 300))) # 维度适配: 截断或填充到decoder的motor_dim dec_motor_dim = self._decoder.motor_dim if motor_vec.size < dec_motor_dim: motor_vec = np.pad(motor_vec, (0, dec_motor_dim - motor_vec.size)) elif motor_vec.size > dec_motor_dim: motor_vec = motor_vec[:dec_motor_dim] # 训练一步: motor向量 + 目标答案 result = self._decoder.train_step(motor_vec, a, lr=lr) epoch_loss += result['loss'] if isinstance(result, dict) else result n_trained += 1 total_loss = epoch_loss # 保存权重 try: dec_path = os.path.join(os.path.dirname(__file__), '..', '..', 'models', 'decoder_weights.npz') self._decoder.save(dec_path) except Exception as e: print(f"[Brain] 解码层权重保存失败: {e}") return {'epochs': epochs, 'n_trained': n_trained, 'final_loss': total_loss} def train(self, texts: list, epochs: int = 1, lr: float = 0.01): """在线训练 — 赫布学习（训练时临时解冻） 关键：必须forward向量而非字符串，否则走chat()分支绕过AreaNetwork """ import time as _t from .semantic_encoder import get_encoder encoder = get_encoder() self._net.unfreeze() try: for epoch in range(epochs): for i, text in enumerate(texts): t0 = _t.time() vec = encoder.encode(text) if isinstance(text, str) else text if vec is None: continue t1 = _t.time() self.forward(vec) t2 = _t.time() self._net.learn(force=True) t3 = _t.time() print(f" train[{i}] encode={t1-t0:.1f}s forward={t2-t1:.1f}s learn={t3-t2:.1f}s", flush=True) finally: self._net.freeze() def fedavg(self, area_name: str, incoming: np.ndarray, node_id: str = '') -> Dict: """联邦平均聚合""" local = self.get_area_weights(area_name) if local.size == 0: self.set_area_weights(area_name, incoming) return {'action': 'adopt', 'node': node_id} # 简单平均 min_len = min(local.size, incoming.size) avg = (local[:min_len] + incoming[:min_len]) / 2 # 补齐 if local.size > min_len: avg = np.concatenate([avg, local[min_len:]]) elif incoming.size > min_len: avg = np.concatenate([avg, incoming[min_len:]]) self.set_area_weights(area_name, avg.astype(np.float32)) return {'action': 'avg', 'local_size': local.size, 'incoming_size': incoming.size} @property def areas(self) -> Dict[str, object]: """兼容: 返回 {区名str: FunctionalArea}""" return {at.value: area for at, area in self._net.areas.items()} def get_stats(self) -> Dict: """兼容旧接口""" return self.stats() def save(self, path: str): """保存模型""" all_weights = {} for area_type, area in self._net.areas.items(): name = area_type.value all_weights[name] = self.get_area_weights(name) np.savez(path, **all_weights) def load(self, path: str): """加载模型""" data = np.load(path, allow_pickle=True) for name in data.files: try: self.set_area_weights(name, data[name]) except (ValueError, KeyError): pass