diff --git a/jspaceai/__init__.py b/jspaceai/__init__.py index e7aba7d..2be3fd0 100644 --- a/jspaceai/__init__.py +++ b/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", ] diff --git a/jspaceai/embodied.py b/jspaceai/embodied.py index eb3276c..6d27d72 100644 --- a/jspaceai/embodied.py +++ b/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 diff --git a/jspaceai/language_model.py b/jspaceai/language_model.py index 434903d..7315ad3 100644 --- a/jspaceai/language_model.py +++ b/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) - self.train() + 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 - self.optimizer = torch.optim.Adam(model.parameters(), lr=lr) + 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: """ diff --git a/jspaceai/training.py b/jspaceai/training.py new file mode 100644 index 0000000..25c7966 --- /dev/null +++ b/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) diff --git a/main.py b/main.py index 83a1ca8..62e9fd5 100644 --- a/main.py +++ b/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) diff --git a/main_chat.py b/main_chat.py index 18b346f..029ee3d 100644 --- a/main_chat.py +++ b/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, ) diff --git a/tests/test_smoke.py b/tests/test_smoke.py index a5ae40a..9939adc 100644 --- a/tests/test_smoke.py +++ b/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() diff --git a/train_chat.py b/train_chat.py index e5535ba..8c2b1a6 100644 --- a/train_chat.py +++ b/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}") + 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', + ) - 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) + 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() - logits, _ = model(toks) - loss = F.cross_entropy( - logits[:, :-1].reshape(-1, cfg.vocab_size), - toks[:, 1:].reshape(-1), - ) - opt.zero_grad() - loss.backward() - torch.nn.utils.clip_grad_norm_(model.parameters(), 0.5) # 更严格的裁剪 - opt.step() + trainer.fit_tokens( + all_tokens, + max_steps=n_steps, + checkpoint_path=mp, + on_progress=report, + ) - 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])}") - - 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")