JspaceAI/jspaceai/evolution.py

"""
自主进化训练器

核心：模型在推理的同时持续学习。每处理一段文本，参数就更新一次。

进化循环：
    1. 喂入新文本片段
    2. forward + 计算 next-token loss
    3. EWC 优化器更新参数（保护旧知识）
    4. 经验回放：当前片段存入 buffer，定期回放旧片段
    5. 周期性 consolidate EWC（更新 Fisher 信息和锚点）
    6. 追踪专家可塑性统计
    7. 定期生成样本，观察进化效果

这个循环可以无限运行——模型永远不会"训练完成"，它一直在进化。
"""
from __future__ import annotations

import torch
import torch.nn.functional as F
from typing import Callable
from .language_model import (
    JSpaceLanguageModel, LanguageConfig,
    ExperienceReplay, EWCOptimizer, ExpertPlasticity,
)


class EvolutionTrainer:
    """自主进化训练器"""

    def __init__(self, model: JSpaceLanguageModel, config: LanguageConfig,
                 lr: float = 1e-3, ewc_lambda: float = 0.1,
                 device: str = 'cpu'):
        self.model = model.to(device)
        self.config = config
        self.device = device

        # 三大自主进化机制
        self.ewc_optimizer = EWCOptimizer(model, lr=lr, ewc_lambda=ewc_lambda)
        self.replay_buffer = ExperienceReplay(capacity=500, seq_len=64)
        self.plasticity = ExpertPlasticity(num_experts=config.num_experts)

        # 进化历史追踪
        self.history: list[dict] = []

    def learn_step(self, token_seq: torch.Tensor) -> dict:
        """单步学习

        Args:
            token_seq: (batch, T) token indices

        Returns:
            stats: 包含 loss、注意力、专家统计等
        """
        token_seq = token_seq.to(self.device)

        # 1. Forward
        logits, info = self.model(token_seq)

        # 2. Next-token prediction loss
        # logits[:, t] 预测 token_seq[:, t+1]
        pred_logits = logits[:, :-1]  # (batch, T-1, vocab)
        targets = token_seq[:, 1:]    # (batch, T-1)
        loss = F.cross_entropy(
            pred_logits.reshape(-1, self.config.vocab_size),
            targets.reshape(-1),
        )

        # 3. 经验回放：如果有足够样本，混入旧数据
        replay_loss = torch.tensor(0.0, device=self.device)
        if len(self.replay_buffer.buffer) >= 8:
            replay_seq = self.replay_buffer.sample(4)
            if replay_seq is not None:
                replay_seq = replay_seq.to(self.device)
                replay_logits, _ = self.model(replay_seq)
                replay_pred = replay_logits[:, :-1]
                replay_targets = replay_seq[:, 1:]
                replay_loss = F.cross_entropy(
                    replay_pred.reshape(-1, self.config.vocab_size),
                    replay_targets.reshape(-1),
                )

        # 4. 总 loss + EWC 正则 + 经验回放
        total_task_loss = loss + 0.5 * replay_loss
        total_loss = self.ewc_optimizer.step(total_task_loss)

        # 5. 更新专家可塑性统计
        self.plasticity.update(info['alpha'].detach(), token_seq.detach())

        # 6. 存入经验回放
        self.replay_buffer.push(token_seq.detach().cpu())

        stats = {
            'loss': loss.item(),
            'replay_loss': replay_loss.item() if isinstance(replay_loss, torch.Tensor) else replay_loss,
            'total_loss': total_loss,
            'w_norm_mean': info['w_norm'].mean().item(),
            'alpha_mean': info['alpha'].mean(dim=(0, 1)).detach().cpu().tolist(),
        }
        return stats

    def consolidate(self, data_sample: torch.Tensor):
        """周期性 consolidate EWC——更新参数重要性锚点

        在学完一段文本后调用，把当前知识"固化"
        """
        self.ewc_optimizer.consolidate(data_sample.to(self.device), n_samples=20)

    def evolve(self, text_stream: list[str], tokenizer,
               seq_len: int = 64, batch_size: int = 4,
               consolidate_every: int = 20,
               generate_every: int = 50,
               max_steps: int | None = None,
               prompt_text: str = "To be",
               on_progress: Callable | None = None) -> list[dict]:
        """
        持续进化主循环

        Args:
            text_stream: 文本片段列表（模拟持续到来的数据流）
            tokenizer: CharTokenizer
            seq_len: 序列长度
            batch_size: 每次喂入的 batch 大小
            consolidate_every: 每隔多少步 consolidate EWC
            generate_every: 每隔多少步生成样本观察
            prompt_text: 生成样本的提示词
            on_progress: 回调函数，返回当前进度

        Returns:
            history: 进化历史
        """
        step = 0
        all_tokens = []

        # 把所有文本编码成 token 流
        for text in text_stream:
            tokens = tokenizer.encode(text)
            all_tokens.extend(tokens)

        # 用滑动窗口在完整 token 流上采样 batch
        # 每个 batch 包含 batch_size 条序列，每条长 seq_len
        # 相邻 batch 之间步进 stride 个序列
        stride = batch_size  # 每个 batch 用 batch_size 个新起点
        n_possible_starts = max(0, len(all_tokens) - seq_len - 1)
        n_batches = max(0, n_possible_starts // stride)

        print(f"自主进化开始：{len(all_tokens)} tokens, {len(text_stream)} 段文本")
        print(f"配置: seq_len={seq_len}, batch_size={batch_size}, "
              f"stride={stride}, n_batches={n_batches}")
        print("=" * 70)

        for batch_idx in range(n_batches):
            # 取一个 batch 的序列（滑动窗口）
            batch_tokens = []
            for i in range(batch_size):
                start = batch_idx * stride + i
                if start + seq_len >= len(all_tokens):
                    batch_tokens.append(all_tokens[-seq_len:])
                else:
                    batch_tokens.append(all_tokens[start:start + seq_len])

            if any(len(b) < seq_len for b in batch_tokens):
                continue

            token_seq = torch.tensor(batch_tokens, dtype=torch.long)
            stats = self.learn_step(token_seq)
            stats['step'] = step
            stats['batch_idx'] = batch_idx
            self.history.append(stats)

            # 周期性 consolidate
            if step > 0 and step % consolidate_every == 0:
                self.consolidate(token_seq)

            # 周期性生成 + 报告
            if step % generate_every == 0 or step == n_batches - 1:
                prompt_ids = tokenizer.encode(prompt_text)
                generated = self.model.generate(
                    prompt_ids, n_new=80, temperature=0.8, top_k=5
                )
                sample = prompt_text + tokenizer.decode(generated)
                stats['sample'] = sample

                print(f"\n[step {step:4d}] loss={stats['loss']:.4f} "
                      f"replay={stats['replay_loss']:.4f} "
                      f"||w||={stats['w_norm_mean']:.3f}")
                print(f"  专家使用率: {[f'{u:.2f}' for u in self.plasticity.usage.tolist()]}")
                print(f"  生成样本: {repr(sample[:120])}...")

                if on_progress:
                    on_progress(stats)

            step += 1

            if max_steps is not None and step >= max_steps:
                break

        print("\n" + "=" * 70)
        print("自主进化完成")
        return self.history

    def get_evolution_summary(self) -> dict:
        """获取进化总结"""
        if not self.history:
            return {}

        losses = [h['loss'] for h in self.history]
        return {
            'steps': len(self.history),
            'final_loss': losses[-1],
            'initial_loss': losses[0],
            'min_loss': min(losses),
            'expert_usage': self.plasticity.usage.tolist(),
            'expert_specialization': self.plasticity.get_stats()['top_specialization'],
            'samples': [h.get('sample', '') for h in self.history if 'sample' in h],
        }