feat: 更新模块结构，添加训练相关功能和可恢复训练会话，优化模型生成逻辑

jspaceai/__init__.py

@@ -3,11 +3,12 @@ JspaceAI —— 全局工作空间 + J-space 广播的智慧系统
 
 模块：
     1. core.py: 核心架构（Expert + JSpaceWorkspace + JSpaceModel，含 RK4/LayerNorm/异构专家）
-    2. language_model.py: 语言版 + 自主进化
+    2. language_model.py: 语言版 JSpace 模型
     3. jlens.py: J-lens 可解释性工具
     4. multimodal.py: 多模态（图像/音频/视频/文本）
-    5. realtime.py: 实时 I/O（摄像头/麦克风/扬声器）
-    6. evolution.py: 自主进化训练器
+    5. policy.py / runtime.py: 动作策略与 workspace 主循环
+    6. events.py / memory.py: 可序列化事件与记忆接口
+    7. training.py / continual.py: 可恢复训练与在线学习
 """
 from .core import (
     Expert,
@@ -59,6 +60,13 @@ from .multimodal import (
 from .continual import (
     OnlineLanguageLearner,
 )
+from .training import (
+    LanguageTrainingConfig,
+    LanguageTrainingSession,
+    TokenBatchSampler,
+    expert_integration_mode,
+    save_language_checkpoint,
+)
 from .policy import (
     ACTION_LABELS,
     compose_action_params,
@@ -156,4 +164,7 @@ __all__ = [
     # 自主进化
     "EvolutionTrainer",
     "OnlineLanguageLearner",
+    "LanguageTrainingConfig", "LanguageTrainingSession",
+    "TokenBatchSampler", "expert_integration_mode",
+    "save_language_checkpoint",
 ]

jspaceai/embodied.py

@@ -1,21 +1,9 @@
 """
-输出执行器层 + 神经系统
+Embodied runtime adapter.
 
-对应人类神经系统的各部分：
-    - 大脑皮层: workspace w + 专家池（已在 multimodal.py）
-    - 小脑: 运动控制器（前向模型+逆模型，精细动作）
-    - 中枢神经: 动作调度器（反射弧+决策门控）
-    - 海马体: 外部情景记忆库
-    - 基底神经节: 动作价值学习（习惯化）
-    - 执行器: 鼠标控制 + 键盘输出 + 音频输出 + 屏幕绘制
-
-核心思想：输出和输入对称。
-    输入：摄像头/麦克风/屏幕/键盘/鼠标 → 编码 → workspace
-    输出：workspace → 解码 → 鼠标移动/键盘按键/音频播放/屏幕绘制
-
-workspace 是模态无关的"意图空间"。
-"想点击左上角"这个意图，在 workspace 里是一个向量，
-解码到鼠标控制器就是移动+点击，解码到键盘就是 Tab+Enter。
+This module owns local sensors/effectors and delegates decision making to
+ActionPolicy. The shared state remains the workspace vector plus serializable
+events, so memory and runtime orchestration can evolve independently.
 """
 from __future__ import annotations
 

jspaceai/language_model.py

@@ -199,6 +199,7 @@ class JSpaceLanguageModel(nn.Module):
             temperature: 采样温度
             top_k: top-k 采样
         """
+        was_training = self.training
         self.eval()
         device = next(self.parameters()).device
         state = self.init_state(1, device)
@@ -226,7 +227,10 @@ class JSpaceLanguageModel(nn.Module):
             generated.append(next_tok)
             tokens.append(next_tok)
 
+        if was_training:
             self.train()
+        else:
+            self.eval()
         return generated
 
 
@@ -285,11 +289,18 @@ class EWCOptimizer:
     """
 
     def __init__(self, model: nn.Module, lr: float = 1e-3,
-                 ewc_lambda: float = 1.0, max_grad_norm: float = 1.0):
+                 ewc_lambda: float = 1.0, max_grad_norm: float = 1.0,
+                 weight_decay: float = 0.0):
         self.model = model
+        if weight_decay > 0:
+            self.optimizer = torch.optim.AdamW(
+                model.parameters(), lr=lr, weight_decay=weight_decay,
+            )
+        else:
             self.optimizer = torch.optim.Adam(model.parameters(), lr=lr)
         self.ewc_lambda = ewc_lambda
         self.max_grad_norm = max_grad_norm
+        self.weight_decay = weight_decay
 
         # Fisher 信息和锚定参数
         self.fisher: dict[str, torch.Tensor] = {}
@@ -361,6 +372,25 @@ class EWCOptimizer:
         self.optimizer.step()
         return total_loss.item()
 
+    def state_dict(self) -> dict:
+        return {
+            "optimizer": self.optimizer.state_dict(),
+            "ewc_lambda": self.ewc_lambda,
+            "max_grad_norm": self.max_grad_norm,
+            "weight_decay": self.weight_decay,
+            "fisher": self.fisher,
+            "anchored_params": self.anchored_params,
+        }
+
+    def load_state_dict(self, state: dict):
+        if "optimizer" in state:
+            self.optimizer.load_state_dict(state["optimizer"])
+        self.ewc_lambda = state.get("ewc_lambda", self.ewc_lambda)
+        self.max_grad_norm = state.get("max_grad_norm", self.max_grad_norm)
+        self.weight_decay = state.get("weight_decay", self.weight_decay)
+        self.fisher = state.get("fisher", {})
+        self.anchored_params = state.get("anchored_params", {})
+
 
 class ExpertPlasticity:
     """

jspaceai/training.py

@@ -0,0 +1,311 @@
+"""
+Reusable training utilities for the JSpace language model.
+
+The goal is to keep training scalable without hard-wiring it to one script:
+sampling, validation, checkpointing, replay, and EWC all live behind a single
+session object that can be reused by CLI tools, tests, and future workers.
+"""
+from __future__ import annotations
+
+from contextlib import contextmanager
+from dataclasses import asdict, dataclass
+from pathlib import Path
+from typing import Callable
+
+import torch
+import torch.nn.functional as F
+
+from .language_data import CharTokenizer
+from .language_model import (
+    EWCOptimizer,
+    ExperienceReplay,
+    ExpertPlasticity,
+    JSpaceLanguageModel,
+    LanguageConfig,
+)
+
+
+@dataclass
+class LanguageTrainingConfig:
+    seq_len: int = 64
+    batch_size: int = 8
+    lr: float = 1e-3
+    weight_decay: float = 0.0
+    ewc_lambda: float = 0.05
+    max_grad_norm: float = 0.5
+    replay_capacity: int = 500
+    replay_batch_size: int = 4
+    replay_weight: float = 0.5
+    consolidate_every: int = 50
+    consolidate_samples: int = 10
+    validate_every: int = 50
+    validate_batches: int = 4
+    save_every: int = 100
+    train_fraction: float = 0.98
+    use_euler_during_train: bool = True
+
+
+@contextmanager
+def expert_integration_mode(model, use_rk4: bool):
+    """Temporarily set expert integration mode."""
+    original = [expert.use_rk4 for expert in model.experts]
+    for expert in model.experts:
+        expert.use_rk4 = use_rk4
+    try:
+        yield
+    finally:
+        for expert, enabled in zip(model.experts, original):
+            expert.use_rk4 = enabled
+
+
+class TokenBatchSampler:
+    """Random contiguous sampler over a token stream with a train/val split."""
+
+    def __init__(
+        self,
+        token_ids: list[int],
+        seq_len: int = 64,
+        train_fraction: float = 0.98,
+    ):
+        if not token_ids:
+            raise ValueError("token_ids must not be empty")
+        self.seq_len = seq_len
+        min_len = seq_len + 2
+        if len(token_ids) < min_len:
+            repeats = min_len // len(token_ids) + 1
+            token_ids = (token_ids * repeats)[:min_len]
+
+        split = int(len(token_ids) * train_fraction)
+        split = min(max(split, min_len), len(token_ids))
+        self.train_tokens = token_ids[:split]
+        self.val_tokens = token_ids[split:] if len(token_ids) - split >= min_len else token_ids[:split]
+
+    def sample(self, batch_size: int, split: str = "train", device: str = "cpu") -> torch.Tensor:
+        tokens = self.train_tokens if split == "train" else self.val_tokens
+        max_start = max(1, len(tokens) - self.seq_len - 1)
+        rows = []
+        for _ in range(batch_size):
+            start = torch.randint(0, max_start, (1,)).item()
+            rows.append(tokens[start:start + self.seq_len])
+        return torch.tensor(rows, dtype=torch.long, device=device)
+
+
+def save_language_checkpoint(
+    path: str | Path,
+    model: JSpaceLanguageModel,
+    config: LanguageConfig,
+    tokenizer: CharTokenizer,
+    trainer_state: dict | None = None,
+    metadata: dict | None = None,
+):
+    path = Path(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    torch.save({
+        "model": model.state_dict(),
+        "config": config,
+        "tokenizer_chars": tokenizer.chars,
+        "trainer": trainer_state or {},
+        "metadata": metadata or {},
+    }, path)
+
+
+class LanguageTrainingSession:
+    """Stateful trainer for scalable language-model training."""
+
+    def __init__(
+        self,
+        model: JSpaceLanguageModel,
+        model_config: LanguageConfig,
+        tokenizer: CharTokenizer,
+        train_config: LanguageTrainingConfig | None = None,
+        device: str = "cpu",
+    ):
+        self.model = model.to(device)
+        self.model_config = model_config
+        self.tokenizer = tokenizer
+        self.train_config = train_config or LanguageTrainingConfig()
+        self.device = device
+        self.optimizer = EWCOptimizer(
+            self.model,
+            lr=self.train_config.lr,
+            ewc_lambda=self.train_config.ewc_lambda,
+            max_grad_norm=self.train_config.max_grad_norm,
+            weight_decay=self.train_config.weight_decay,
+        )
+        self.replay_buffer = ExperienceReplay(
+            capacity=self.train_config.replay_capacity,
+            seq_len=self.train_config.seq_len,
+        )
+        self.plasticity = ExpertPlasticity(num_experts=model_config.num_experts)
+        self.global_step = 0
+        self.history: list[dict] = []
+
+    def learn_batch(self, token_seq: torch.Tensor) -> dict:
+        token_seq = token_seq.to(self.device)
+        self.model.train()
+        logits, info = self.model(token_seq)
+        pred = logits[:, :-1]
+        target = token_seq[:, 1:]
+        loss = F.cross_entropy(
+            pred.reshape(-1, self.model_config.vocab_size),
+            target.reshape(-1),
+        )
+
+        replay_loss = torch.tensor(0.0, device=self.device)
+        replay_seq = self.replay_buffer.sample(self.train_config.replay_batch_size)
+        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_target = replay_seq[:, 1:]
+            replay_loss = F.cross_entropy(
+                replay_pred.reshape(-1, self.model_config.vocab_size),
+                replay_target.reshape(-1),
+            )
+
+        total_task_loss = loss + self.train_config.replay_weight * replay_loss
+        total_loss = self.optimizer.step(total_task_loss)
+
+        self.plasticity.update(info["alpha"].detach(), token_seq.detach())
+        self.replay_buffer.push(token_seq.detach().cpu())
+
+        return {
+            "loss": float(loss.item()),
+            "replay_loss": float(replay_loss.item()),
+            "total_loss": float(total_loss),
+            "w_norm_mean": float(info["w_norm"].mean().item()),
+            "expert_usage": self.plasticity.usage.tolist(),
+        }
+
+    @torch.no_grad()
+    def evaluate(self, sampler: TokenBatchSampler) -> float:
+        was_training = self.model.training
+        self.model.eval()
+        losses = []
+        for _ in range(max(1, self.train_config.validate_batches)):
+            token_seq = sampler.sample(
+                self.train_config.batch_size,
+                split="val",
+                device=self.device,
+            )
+            logits, _ = self.model(token_seq)
+            loss = F.cross_entropy(
+                logits[:, :-1].reshape(-1, self.model_config.vocab_size),
+                token_seq[:, 1:].reshape(-1),
+            )
+            losses.append(loss.item())
+        if was_training:
+            self.model.train()
+        return float(sum(losses) / len(losses))
+
+    def fit_text(
+        self,
+        text: str,
+        max_steps: int,
+        checkpoint_path: str | Path | None = None,
+        on_progress: Callable[[dict], None] | None = None,
+    ) -> list[dict]:
+        return self.fit_tokens(
+            self.tokenizer.encode(text),
+            max_steps=max_steps,
+            checkpoint_path=checkpoint_path,
+            on_progress=on_progress,
+        )
+
+    def fit_tokens(
+        self,
+        token_ids: list[int],
+        max_steps: int,
+        checkpoint_path: str | Path | None = None,
+        on_progress: Callable[[dict], None] | None = None,
+    ) -> list[dict]:
+        sampler = TokenBatchSampler(
+            token_ids,
+            seq_len=self.train_config.seq_len,
+            train_fraction=self.train_config.train_fraction,
+        )
+        use_rk4 = not self.train_config.use_euler_during_train
+        with expert_integration_mode(self.model, use_rk4=use_rk4):
+            for _ in range(max_steps):
+                batch = sampler.sample(
+                    self.train_config.batch_size,
+                    split="train",
+                    device=self.device,
+                )
+                stats = self.learn_batch(batch)
+                self.global_step += 1
+                stats["step"] = self.global_step
+
+                if (
+                    self.train_config.validate_every > 0
+                    and self.global_step % self.train_config.validate_every == 0
+                ):
+                    stats["val_loss"] = self.evaluate(sampler)
+
+                if (
+                    self.train_config.consolidate_every > 0
+                    and self.global_step % self.train_config.consolidate_every == 0
+                ):
+                    self.optimizer.consolidate(
+                        batch.detach(),
+                        n_samples=self.train_config.consolidate_samples,
+                    )
+
+                self.history.append(stats)
+                if on_progress:
+                    on_progress(stats)
+
+                if (
+                    checkpoint_path is not None
+                    and self.train_config.save_every > 0
+                    and self.global_step % self.train_config.save_every == 0
+                ):
+                    self.save_checkpoint(checkpoint_path)
+
+        if checkpoint_path is not None:
+            self.save_checkpoint(checkpoint_path)
+        return self.history
+
+    def state_dict(self) -> dict:
+        return {
+            "global_step": self.global_step,
+            "optimizer": self.optimizer.state_dict(),
+            "replay_buffer": list(self.replay_buffer.buffer),
+            "plasticity": {
+                "usage": self.plasticity.usage,
+                "expert_specialization": self.plasticity.expert_specialization,
+            },
+            "train_config": asdict(self.train_config),
+            "history": self.history[-200:],
+        }
+
+    def load_state_dict(self, state: dict):
+        self.global_step = int(state.get("global_step", 0))
+        if "optimizer" in state:
+            self.optimizer.load_state_dict(state["optimizer"])
+        replay = state.get("replay_buffer", [])
+        self.replay_buffer.buffer.clear()
+        for seq in replay:
+            self.replay_buffer.push(seq)
+        plasticity = state.get("plasticity", {})
+        if "usage" in plasticity:
+            self.plasticity.usage = plasticity["usage"].detach().cpu()
+        if "expert_specialization" in plasticity:
+            self.plasticity.expert_specialization = plasticity["expert_specialization"]
+        self.history = list(state.get("history", []))
+
+    def save_checkpoint(self, path: str | Path, metadata: dict | None = None):
+        save_language_checkpoint(
+            path,
+            self.model,
+            self.model_config,
+            self.tokenizer,
+            trainer_state=self.state_dict(),
+            metadata=metadata,
+        )
+
+    def load_checkpoint_state(self, path: str | Path):
+        ckpt = torch.load(path, map_location=self.device, weights_only=False)
+        trainer_state = ckpt.get("trainer")
+        if trainer_state:
+            self.load_state_dict(trainer_state)

main.py

@@ -209,7 +209,7 @@ def live(n_steps: int, device: str, safe_mode: bool = False, unsafe: bool = Fals
         # 记忆数
         ax = axes[1, 2]
         ax.plot(steps, [s['memory_count'] for s in log], 'teal')
-        ax.set_title('Hippocampus Memory Count'); ax.grid(True, alpha=0.3)
+        ax.set_title('Memory Count'); ax.grid(True, alpha=0.3)
 
         plt.tight_layout()
         Path('outputs').mkdir(exist_ok=True)

main_chat.py

@@ -7,7 +7,7 @@ JspaceAI —— 对话版（只有语言，控制台交互）
 
 模式：
     --mode chat:     交互对话（默认）
-    --mode train:    先在 Shakespeare 语料上预训练若干步，再进入对话
+    --mode train:    在清洗后的中文语料上训练若干步
     --mode generate: 给定提示词一次性生成
 
 运行：
@@ -21,9 +21,11 @@ import torch
 from pathlib import Path
 
 from jspaceai import (
-    LanguageConfig, JSpaceLanguageModel, EvolutionTrainer,
+    LanguageConfig, JSpaceLanguageModel,
     OnlineLanguageLearner,
-    CharTokenizer, load_chinese_corpus,
+    CharTokenizer,
+    LanguageTrainingConfig, LanguageTrainingSession,
+    expert_integration_mode, save_language_checkpoint,
 )
 from train_chat import clean_corpus
 
@@ -80,12 +82,7 @@ def load_or_init_model(device: str):
 
 
 def save_model(model, config, tokenizer):
-    Path('outputs').mkdir(exist_ok=True)
-    torch.save({
-        'model': model.state_dict(),
-        'config': config,
-        'tokenizer_chars': tokenizer.chars,
-    }, 'outputs/chat_model.pt')
+    save_language_checkpoint('outputs/chat_model.pt', model, config, tokenizer)
 
 
 def ensure_corpus(text: str | None) -> str:
@@ -95,59 +92,45 @@ def ensure_corpus(text: str | None) -> str:
     return text
 
 
-class temporary_rk4:
-    """Temporarily switch expert integration mode for faster inference."""
-
-    def __init__(self, model, enabled: bool):
-        self.model = model
-        self.enabled = enabled
-        self.original = []
-
-    def __enter__(self):
-        self.original = [expert.use_rk4 for expert in self.model.experts]
-        for expert in self.model.experts:
-            expert.use_rk4 = self.enabled
-
-    def __exit__(self, exc_type, exc, tb):
-        for expert, enabled in zip(self.model.experts, self.original):
-            expert.use_rk4 = enabled
-        return False
-
-
 def train(model, tokenizer, text: str | None, n_steps: int, device: str):
-    """在 Shakespeare 语料上预训练
-
-    训练时临时关闭 RK4 用 Euler 加速（快 4 倍），训练完恢复 RK4。
-    """
+    """Scalable language-model training entrypoint."""
     print("\n" + "=" * 60)
-    print(f"预训练 {n_steps} 步（Shakespeare 语料）")
+    print(f"训练 {n_steps} 步")
     print("=" * 60)
 
     text = ensure_corpus(text)
     config = model.config
-
-    # 训练时临时关 RK4 加速（Euler 快 4 倍）
-    original_rk4 = config.use_rk4
-    for expert in model.experts:
-        expert.use_rk4 = False
-    print(f"训练模式: Euler（加速），训练后恢复 RK4")
-
-    trainer = EvolutionTrainer(
-        model, config, lr=5e-3, ewc_lambda=0.05, device=device,
+    train_cfg = LanguageTrainingConfig(
+        seq_len=64,
+        batch_size=8,
+        lr=5e-3,
+        ewc_lambda=0.05,
+        consolidate_every=50,
+        validate_every=max(1, min(50, n_steps)),
+        save_every=max(1, min(100, n_steps)),
+        use_euler_during_train=True,
     )
-    chunks = [text[i:i+200] for i in range(0, len(text), 200)]
-    trainer.evolve(
-        chunks, tokenizer,
-        seq_len=64, batch_size=8,
-        consolidate_every=50, generate_every=50,
-        max_steps=n_steps, prompt_text="学而时习之",
+    trainer = LanguageTrainingSession(
+        model, config, tokenizer, train_cfg, device=device,
     )
 
-    # 恢复 RK4
-    for expert in model.experts:
-        expert.use_rk4 = original_rk4
+    def report(stats: dict):
+        step = stats["step"]
+        interval = max(1, min(50, n_steps))
+        if step == 1 or step % interval == 0:
+            val = f" val={stats['val_loss']:.3f}" if "val_loss" in stats else ""
+            print(
+                f"  step {step:4d} | loss={stats['loss']:.3f} "
+                f"replay={stats['replay_loss']:.3f}{val} "
+                f"||w||={stats['w_norm_mean']:.3f}"
+            )
 
-    save_model(model, config, tokenizer)
+    trainer.fit_text(
+        text,
+        max_steps=n_steps,
+        checkpoint_path='outputs/chat_model.pt',
+        on_progress=report,
+    )
     print(f"\n模型已保存: outputs/chat_model.pt")
 
 
@@ -160,7 +143,7 @@ def generate_response(model, tokenizer, prompt: str, n_new: int = 60,
     if not prompt_ids:
         prompt_ids = [0]
     was_training = model.training
-    with temporary_rk4(model, enabled=not fast):
+    with expert_integration_mode(model, use_rk4=not fast):
         generated = model.generate(
             prompt_ids, n_new=n_new, temperature=temperature, top_k=top_k,
         )

tests/test_smoke.py

@@ -14,14 +14,16 @@ from jspaceai import (
     JSpaceLanguageModel,
     JSpaceModel,
     LanguageConfig,
+    LanguageTrainingConfig,
+    LanguageTrainingSession,
     MultimodalConfig,
     MultimodalJSpaceModel,
     OnlineLanguageLearner,
     WorkspaceEvent,
     WorkspaceRuntime,
+    compose_action_params,
 )
 from main_chat import generate_response
-from jspaceai import compose_action_params
 
 
 class SmokeTests(unittest.TestCase):
@@ -75,6 +77,26 @@ class SmokeTests(unittest.TestCase):
         self.assertFalse(model.training)
         self.assertTrue(all(expert.use_rk4 for expert in model.experts))
 
+    def test_language_generate_preserves_training_mode(self):
+        tokenizer = CharTokenizer.from_text("abcabc")
+        config = LanguageConfig(
+            vocab_size=tokenizer.vocab_size,
+            embed_dim=8,
+            input_dim=8,
+            workspace_dim=16,
+            expert_dim=8,
+            num_experts=3,
+            ode_steps=1,
+            noise_std=0.0,
+        )
+        model = JSpaceLanguageModel(config)
+        model.eval()
+        model.generate([1], n_new=1)
+        self.assertFalse(model.training)
+        model.train()
+        model.generate([1], n_new=1)
+        self.assertTrue(model.training)
+
     def test_multimodal_single_step_records_trajectory(self):
         config = MultimodalConfig(
             vocab_size=20,
@@ -149,6 +171,46 @@ class SmokeTests(unittest.TestCase):
         self.assertEqual(second["step"], 2)
         self.assertGreaterEqual(second["replay_loss"], 0.0)
 
+    def test_language_training_session_saves_checkpoint(self):
+        tokenizer = CharTokenizer.from_text("学而时习之学而时习之")
+        config = LanguageConfig(
+            vocab_size=tokenizer.vocab_size,
+            embed_dim=8,
+            input_dim=8,
+            workspace_dim=16,
+            expert_dim=8,
+            num_experts=3,
+            ode_steps=1,
+            noise_std=0.0,
+        )
+        model = JSpaceLanguageModel(config)
+        train_config = LanguageTrainingConfig(
+            seq_len=8,
+            batch_size=2,
+            lr=1e-2,
+            replay_batch_size=1,
+            validate_every=1,
+            validate_batches=1,
+            save_every=1,
+            consolidate_every=0,
+        )
+        session = LanguageTrainingSession(
+            model, config, tokenizer, train_config, device="cpu",
+        )
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            ckpt_path = Path(tmpdir) / "model.pt"
+            history = session.fit_text(
+                "学而时习之学而时习之",
+                max_steps=2,
+                checkpoint_path=ckpt_path,
+            )
+            ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)
+
+        self.assertEqual(len(history), 2)
+        self.assertIn("trainer", ckpt)
+        self.assertEqual(ckpt["trainer"]["global_step"], 2)
+
     def test_compose_action_params_respects_discrete_mode(self):
         raw = torch.tensor([0.5, -0.25, 0.3, -0.4, 0.8]).numpy()
 

train_chat.py

@@ -1,22 +1,23 @@
 #!/usr/bin/env python3
-"""训练对话模型——清洗语料 + 分阶段训练
+"""训练对话模型——清洗语料 + 可恢复分阶段训练
 
 策略：
-  1. 从 corpus/*.txt 加载维基百科语料，繁简转换
-  2. 清洗：去掉 ## 标题行、# 章节标记、过短段落、纯英文段落
-  3. 只保留连续中文段落（>= 20 字符），拼成语料
-  4. vocab=400，只保留高频字符
-  5. 分阶段训练：lr 2e-3 → 1e-3 → 5e-4
+  1. 递归读取 corpus/ 下的本地语料
+  2. 清理 markdown/HTML 噪声，过滤过短或过长段落
+  3. 同时投喂简体和繁体版本
+  4. 统一使用 LanguageTrainingSession 训练、验证和保存
+  5. 分阶段降低学习率
 """
 import torch
-import torch.nn.functional as F
 from pathlib import Path
 import time
 import re
+import os
 
 from jspaceai import (
     LanguageConfig, JSpaceLanguageModel,
-    CharTokenizer, load_chinese_corpus, load_textbook_corpus,
+    CharTokenizer, load_chinese_corpus,
+    LanguageTrainingConfig, LanguageTrainingSession,
 )
 
 
@@ -29,22 +30,42 @@ def get_config(vocab_size: int) -> LanguageConfig:
     )
 
 
-def clean_corpus() -> str:
+def _corpus_budget_mb(default: int | None = 64) -> int | None:
+    raw = os.environ.get("JSPACE_CORPUS_MAX_MB")
+    if raw is None:
+        return default
+    raw = raw.strip().lower()
+    if raw in {"", "0", "all", "none"}:
+        return None
+    return int(raw)
+
+
+def clean_corpus(max_source_mb: int | None = None, use_cache: bool = True) -> str:
     """构建中文训练语料——读取 corpus/ 下所有文件，繁简双版本同时投喂。
 
     策略：
-      1. 递归读取 corpus/ 目录下所有文件（.txt .md 等）
+      1. 递归读取 corpus/ 目录下的本地语料（默认有读取预算）
       2. 去掉 markdown/HTML 标记（链接、表格、标题符号等），只保留纯文本
       3. 对每段文本同时生成简体版和繁体版，都喂给模型
       4. 加上内嵌唐诗宋词论语（简体连续文本）
     """
     from opencc import OpenCC
-    from pathlib import Path
+
+    if max_source_mb is None:
+        max_source_mb = _corpus_budget_mb()
+
+    cache_key = "all" if max_source_mb is None else f"{max_source_mb}mb"
+    cache_path = Path("outputs/cache") / f"clean_corpus_{cache_key}.txt"
+    if use_cache and cache_path.exists():
+        return cache_path.read_text(encoding="utf-8")
+
     cc_s2t = OpenCC('s2t')  # 简转繁
     cc_t2s = OpenCC('t2s')  # 繁转简
 
     corpus_dir = Path(__file__).parent / 'corpus'
     raw_paragraphs = []
+    source_budget = None if max_source_mb is None else max_source_mb * 1024 * 1024
+    bytes_read = 0
 
     # 1. 内嵌唐诗宋词论语
     for para in load_chinese_corpus().split('\n\n'):
@@ -66,10 +87,20 @@ def clean_corpus() -> str:
         for filepath in sorted(corpus_dir.rglob('*')):
             if not filepath.is_file():
                 continue
+            if source_budget is not None:
+                try:
+                    file_size = filepath.stat().st_size
+                except OSError:
+                    continue
+                if bytes_read >= source_budget:
+                    break
+                if file_size > max(1024 * 1024, source_budget - bytes_read):
+                    continue
             try:
                 content = filepath.read_text(encoding='utf-8', errors='ignore')
             except Exception:
                 continue
+            bytes_read += len(content.encode('utf-8', errors='ignore'))
             # 去 markdown 标记
             content = re.sub(r'\[([^\]]*)\]\([^)]*\)', r'\1', content)  # 链接
             content = re.sub(r'^#{1,6}\s+', '', content, flags=re.M)    # 标题
@@ -107,29 +138,11 @@ def clean_corpus() -> str:
         bilingual.append(simplified)
         bilingual.append(traditional)
 
-    return '\n\n'.join(bilingual)
-
-
-def gen_sample(model, tok, prompt_text, n_new=80, temp=0.7, top_k=5):
-    model.eval()
-    with torch.no_grad():
-        prompt = tok.encode(prompt_text)
-        if not prompt:
-            prompt = [0]
-        state = model.init_state(1, torch.device('cpu'))
-        for t in prompt:
-            state, _, _, _ = model.step(state, torch.tensor([t]))
-        gen = []
-        last = prompt[-1]
-        for _ in range(n_new):
-            state, logits, _, _ = model.step(state, torch.tensor([last]))
-            probs = F.softmax(logits[0] / temp, dim=-1)
-            topk = probs.topk(top_k)
-            next_tok = topk.indices[torch.multinomial(topk.values, 1)].item()
-            gen.append(next_tok)
-            last = next_tok
-    model.train()
-    return prompt_text + tok.decode(gen)
+    result = '\n\n'.join(bilingual)
+    if use_cache:
+        cache_path.parent.mkdir(parents=True, exist_ok=True)
+        cache_path.write_text(result, encoding="utf-8")
+    return result
 
 
 def main():
@@ -141,11 +154,7 @@ def main():
     model = JSpaceLanguageModel(cfg)
     print(f"vocab={tok.vocab_size}, params={sum(p.numel() for p in model.parameters()):,}")
 
-    for e in model.experts:
-        e.use_rk4 = False  # Euler 加速
-
     all_tokens = tok.encode(text)
-    seq_len = 64
 
     # 加载已有模型
     mp = Path('outputs/chat_model.pt')
@@ -169,45 +178,41 @@ def main():
 
     t_total = time.time()
     for lr, n_steps, label in stages:
-        opt = torch.optim.Adam(model.parameters(), lr=lr)
         print(f"\n{'='*60}")
         print(f"阶段: {label} | lr={lr} | {n_steps}步")
         print(f"{'='*60}")
-
-        for step in range(n_steps):
-            batch = []
-            for _ in range(8):
-                start = torch.randint(0, max(1, len(all_tokens) - seq_len - 1), (1,)).item()
-                batch.append(all_tokens[start:start + seq_len])
-            toks = torch.tensor(batch, dtype=torch.long)
-
-            logits, _ = model(toks)
-            loss = F.cross_entropy(
-                logits[:, :-1].reshape(-1, cfg.vocab_size),
-                toks[:, 1:].reshape(-1),
+        train_cfg = LanguageTrainingConfig(
+            seq_len=64,
+            batch_size=8,
+            lr=lr,
+            ewc_lambda=0.05,
+            consolidate_every=100,
+            validate_every=100,
+            save_every=100,
+            use_euler_during_train=True,
+        )
+        trainer = LanguageTrainingSession(
+            model, cfg, tok, train_cfg, device='cpu',
         )
-            opt.zero_grad()
-            loss.backward()
-            torch.nn.utils.clip_grad_norm_(model.parameters(), 0.5)  # 更严格的裁剪
-            opt.step()
 
-            if (step + 1) % 100 == 0:
-                samples = []
-                for p in prompts:
-                    s = gen_sample(model, tok, p, n_new=40, temp=0.7)
-                    samples.append(s)
-                print(f"\n  step {step+1:3d} | loss {loss.item():.3f}")
-                for p, s in zip(prompts, samples):
-                    print(f"    [{p}] {repr(s[:60])}")
+        def report(stats: dict):
+            if stats['step'] % 100 != 0:
+                return
+            val = f" val={stats['val_loss']:.3f}" if "val_loss" in stats else ""
+            print(f"\n  step {stats['step']:3d} | loss {stats['loss']:.3f}{val}")
+            model.eval()
+            for prompt in prompts:
+                generated = model.generate(tok.encode(prompt) or [0], n_new=40, temperature=0.7, top_k=5)
+                print(f"    [{prompt}] {repr((prompt + tok.decode(generated))[:60])}")
+            model.train()
+
+        trainer.fit_tokens(
+            all_tokens,
+            max_steps=n_steps,
+            checkpoint_path=mp,
+            on_progress=report,
+        )
 
-    for e in model.experts:
-        e.use_rk4 = True
-    Path('outputs').mkdir(exist_ok=True)
-    torch.save({
-        'model': model.state_dict(),
-        'config': cfg,
-        'tokenizer_chars': tok.chars,
-    }, 'outputs/chat_model.pt')
     print(f"\n完成，总耗时 {time.time()-t_total:.0f}s，已保存到 outputs/chat_model.pt")