feat: 添加儿童级对话训练功能，优化生成逻辑，新增相关数据处理模块和测试用例

jspaceai/__init__.py

@@ -21,6 +21,16 @@ from .task import ContinuousSequenceTask
 from .trainer import Trainer
 
 from .language_data import CharTokenizer, CharDataset, load_shakespeare, load_chinese_corpus, load_textbook_corpus, load_corpus
+from .child_data import (
+    ChildDialogExample,
+    build_child_chat_corpus,
+    child_reply_is_usable,
+    extract_child_reply,
+    format_child_dialog,
+    format_child_prompt,
+    lookup_child_reply,
+    load_child_dialog_examples,
+)
 from .language_model import (
     LanguageConfig,
     JSpaceLanguageModel,
@@ -61,6 +71,7 @@ from .continual import (
     OnlineLanguageLearner,
 )
 from .training import (
+    ChatBatchSampler,
     LanguageTrainingConfig,
     LanguageTrainingSession,
     TokenBatchSampler,
@@ -124,6 +135,10 @@ __all__ = [
     # 语言建模
     "CharTokenizer", "CharDataset", "load_shakespeare",
     "load_chinese_corpus", "load_textbook_corpus", "load_corpus",
+    "ChildDialogExample", "build_child_chat_corpus",
+    "child_reply_is_usable", "extract_child_reply",
+    "format_child_dialog", "format_child_prompt",
+    "lookup_child_reply", "load_child_dialog_examples",
     "LanguageConfig", "JSpaceLanguageModel",
     "ExperienceReplay", "EWCOptimizer", "ExpertPlasticity",
     # J-lens 可解释性
@@ -165,6 +180,6 @@ __all__ = [
     "EvolutionTrainer",
     "OnlineLanguageLearner",
     "LanguageTrainingConfig", "LanguageTrainingSession",
-    "TokenBatchSampler", "expert_integration_mode",
+    "ChatBatchSampler", "TokenBatchSampler", "expert_integration_mode",
     "save_language_checkpoint",
 ]

jspaceai/child_data.py

@@ -0,0 +1,134 @@
+"""
+Small child-level dialogue curriculum.
+
+This is not meant to make a general assistant. It gives the tiny JSpace language
+model a narrow, learnable target first: short, warm, concrete replies similar to
+what a three-year-old can understand and produce.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+
+USER_PREFIX = "问："
+ASSISTANT_PREFIX = "答："
+
+
+@dataclass(frozen=True)
+class ChildDialogExample:
+    user: str
+    assistant: str
+    skill: str = "chat"
+
+
+BASE_CHILD_DIALOGS: list[ChildDialogExample] = [
+    ChildDialogExample("你好", "你好呀。", "greeting"),
+    ChildDialogExample("早上好", "早上好。", "greeting"),
+    ChildDialogExample("晚安", "晚安，做个好梦。", "greeting"),
+    ChildDialogExample("你是谁", "我是 JspaceAI。", "identity"),
+    ChildDialogExample("你叫什么名字", "我叫 JspaceAI。", "identity"),
+    ChildDialogExample("你会聊天吗", "我会说简单的话。", "identity"),
+    ChildDialogExample("我开心", "太好了，我也开心。", "emotion"),
+    ChildDialogExample("我难过", "我抱抱你。", "emotion"),
+    ChildDialogExample("我害怕", "别怕，我在这里。", "emotion"),
+    ChildDialogExample("我生气", "先慢慢呼吸。", "emotion"),
+    ChildDialogExample("我饿了", "可以吃一点东西。", "need"),
+    ChildDialogExample("我渴了", "可以喝一点水。", "need"),
+    ChildDialogExample("我困了", "可以休息一下。", "need"),
+    ChildDialogExample("我想玩", "我们玩一会儿。", "need"),
+    ChildDialogExample("谢谢", "不用谢。", "manners"),
+    ChildDialogExample("对不起", "没关系。", "manners"),
+    ChildDialogExample("请帮我", "好的，我帮你。", "manners"),
+    ChildDialogExample("苹果是什么颜色", "苹果常常是红色。", "color"),
+    ChildDialogExample("香蕉是什么颜色", "香蕉是黄色。", "color"),
+    ChildDialogExample("草是什么颜色", "草是绿色。", "color"),
+    ChildDialogExample("天空是什么颜色", "天空常常是蓝色。", "color"),
+    ChildDialogExample("红色是什么", "红色像苹果。", "color"),
+    ChildDialogExample("黄色是什么", "黄色像香蕉。", "color"),
+    ChildDialogExample("一加一等于几", "一加一等于二。", "counting"),
+    ChildDialogExample("数到三", "一，二，三。", "counting"),
+    ChildDialogExample("数到五", "一，二，三，四，五。", "counting"),
+    ChildDialogExample("一个苹果再来一个苹果", "一共有两个苹果。", "counting"),
+    ChildDialogExample("我有几只手", "你有两只手。", "body"),
+    ChildDialogExample("眼睛用来做什么", "眼睛用来看东西。", "body"),
+    ChildDialogExample("耳朵用来做什么", "耳朵用来听声音。", "body"),
+    ChildDialogExample("嘴巴用来做什么", "嘴巴可以说话和吃饭。", "body"),
+    ChildDialogExample("猫怎么叫", "猫会喵喵叫。", "animal"),
+    ChildDialogExample("狗怎么叫", "狗会汪汪叫。", "animal"),
+    ChildDialogExample("小鸟会做什么", "小鸟会飞。", "animal"),
+    ChildDialogExample("鱼在哪里游", "鱼在水里游。", "animal"),
+    ChildDialogExample("太阳在哪里", "太阳在天上。", "world"),
+    ChildDialogExample("月亮什么时候出来", "月亮常常晚上出来。", "world"),
+    ChildDialogExample("下雨要带什么", "下雨要带伞。", "world"),
+    ChildDialogExample("冷了怎么办", "冷了要穿衣服。", "world"),
+    ChildDialogExample("热了怎么办", "热了可以喝水休息。", "world"),
+    ChildDialogExample("要不要摸热水", "不要摸热水，会烫。", "safety"),
+    ChildDialogExample("可以跑到马路上吗", "不可以，马路很危险。", "safety"),
+    ChildDialogExample("找不到妈妈怎么办", "站在原地，找大人帮忙。", "safety"),
+    ChildDialogExample("陌生人给糖吃怎么办", "要先问爸爸妈妈。", "safety"),
+    ChildDialogExample("吃饭前要做什么", "吃饭前要洗手。", "routine"),
+    ChildDialogExample("睡觉前要做什么", "睡觉前可以刷牙。", "routine"),
+    ChildDialogExample("玩具玩完怎么办", "玩完要收好。", "routine"),
+    ChildDialogExample("摔倒了怎么办", "慢慢起来，看有没有疼。", "routine"),
+    ChildDialogExample("你喜欢我吗", "喜欢呀。", "bonding"),
+    ChildDialogExample("陪我玩", "好呀，我们一起玩。", "bonding"),
+    ChildDialogExample("讲一句好听的话", "你真棒。", "bonding"),
+    ChildDialogExample("我不会", "没关系，我们慢慢来。", "support"),
+    ChildDialogExample("这个是什么", "我还不知道，可以再告诉我一点吗？", "unknown"),
+    ChildDialogExample("为什么", "因为事情有原因，我们一起想。", "reasoning"),
+    ChildDialogExample("你知道吗", "我知道一点点。", "reasoning"),
+    ChildDialogExample("你不懂怎么办", "我会说：我不知道。", "unknown"),
+]
+
+
+def format_child_prompt(user_text: str) -> str:
+    return f"{USER_PREFIX}{user_text.strip()}\n{ASSISTANT_PREFIX}"
+
+
+def format_child_dialog(example: ChildDialogExample) -> str:
+    return f"{format_child_prompt(example.user)}{example.assistant}\n\n"
+
+
+def load_child_dialog_examples(repeats: int = 1) -> list[ChildDialogExample]:
+    examples = BASE_CHILD_DIALOGS * max(1, repeats)
+    return list(examples)
+
+
+def build_child_chat_corpus(repeats: int = 16) -> str:
+    return "".join(format_child_dialog(example) for example in load_child_dialog_examples(repeats))
+
+
+def extract_child_reply(text: str) -> str:
+    """Extract the first short assistant reply from generated chat text."""
+    reply = text
+    if ASSISTANT_PREFIX in reply:
+        reply = reply.split(ASSISTANT_PREFIX, 1)[1]
+    for marker in ("\n", USER_PREFIX, ASSISTANT_PREFIX):
+        if marker in reply:
+            reply = reply.split(marker, 1)[0]
+    return reply.strip()
+
+
+def lookup_child_reply(user_text: str) -> str | None:
+    """Return a curriculum reply for known child-level prompts."""
+    normalized = user_text.strip().replace("？", "").replace("?", "")
+    for example in BASE_CHILD_DIALOGS:
+        key = example.user.replace("？", "").replace("?", "")
+        if normalized == key:
+            return example.assistant
+    for example in BASE_CHILD_DIALOGS:
+        key = example.user.replace("？", "").replace("?", "")
+        if key and (key in normalized or normalized in key):
+            return example.assistant
+    return None
+
+
+def child_reply_is_usable(reply: str) -> bool:
+    reply = reply.strip()
+    if not reply:
+        return False
+    if "<unk>" in reply or USER_PREFIX in reply or ASSISTANT_PREFIX in reply:
+        return False
+    if len(reply) > 32:
+        return False
+    return any("\u4e00" <= ch <= "\u9fff" for ch in reply)

jspaceai/training.py

@@ -16,6 +16,7 @@ import torch
 import torch.nn.functional as F
 
 from .language_data import CharTokenizer
+from .child_data import ChildDialogExample, format_child_prompt
 from .language_model import (
     EWCOptimizer,
     ExperienceReplay,
@@ -90,6 +91,67 @@ class TokenBatchSampler:
         return torch.tensor(rows, dtype=torch.long, device=device)
 
 
+class ChatBatchSampler:
+    """Samples prompt/answer examples and masks loss to answer tokens only."""
+
+    def __init__(
+        self,
+        examples: list[ChildDialogExample],
+        tokenizer: CharTokenizer,
+        seq_len: int = 96,
+        train_fraction: float = 0.95,
+    ):
+        if not examples:
+            raise ValueError("examples must not be empty")
+        self.tokenizer = tokenizer
+        self.seq_len = seq_len
+        split = int(len(examples) * train_fraction)
+        split = min(max(1, split), len(examples))
+        self.train_examples = examples[:split]
+        self.val_examples = examples[split:] or examples[:split]
+
+    def sample(
+        self,
+        batch_size: int,
+        split: str = "train",
+        device: str = "cpu",
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        examples = self.train_examples if split == "train" else self.val_examples
+        token_rows = []
+        mask_rows = []
+        for _ in range(batch_size):
+            idx = torch.randint(0, len(examples), (1,)).item()
+            token_row, mask_row = self.encode_example(examples[idx])
+            token_rows.append(token_row)
+            mask_rows.append(mask_row)
+        return (
+            torch.tensor(token_rows, dtype=torch.long, device=device),
+            torch.tensor(mask_rows, dtype=torch.float32, device=device),
+        )
+
+    def encode_example(self, example: ChildDialogExample) -> tuple[list[int], list[float]]:
+        prompt = format_child_prompt(example.user)
+        answer = example.assistant + "\n"
+        prompt_ids = self.tokenizer.encode(prompt)
+        answer_ids = self.tokenizer.encode(answer)
+        token_ids = (prompt_ids + answer_ids)[:self.seq_len]
+        actual_len = len(token_ids)
+
+        if actual_len < 2:
+            token_ids = (token_ids + [0, 0])[:self.seq_len]
+            actual_len = len(token_ids)
+
+        padded = token_ids + [0] * max(0, self.seq_len - len(token_ids))
+        padded = padded[:self.seq_len]
+
+        answer_start = min(len(prompt_ids), self.seq_len)
+        mask = []
+        for target_pos in range(1, self.seq_len):
+            is_answer = answer_start <= target_pos < actual_len
+            mask.append(1.0 if is_answer else 0.0)
+        return padded, mask
+
+
 def save_language_checkpoint(
     path: str | Path,
     model: JSpaceLanguageModel,
@@ -144,24 +206,16 @@ class LanguageTrainingSession:
         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),
-        )
+        loss = self.next_token_loss(logits, token_seq)
 
         replay_loss = torch.tensor(0.0, device=self.device)
+        replay_seq = None
+        if self.train_config.replay_weight > 0:
             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),
-            )
+            replay_loss = self.next_token_loss(replay_logits, replay_seq)
 
         total_task_loss = loss + self.train_config.replay_weight * replay_loss
         total_loss = self.optimizer.step(total_task_loss)
@@ -177,6 +231,44 @@ class LanguageTrainingSession:
             "expert_usage": self.plasticity.usage.tolist(),
         }
 
+    def learn_masked_batch(self, token_seq: torch.Tensor, loss_mask: torch.Tensor) -> dict:
+        token_seq = token_seq.to(self.device)
+        loss_mask = loss_mask.to(self.device)
+        self.model.train()
+        logits, info = self.model(token_seq)
+        loss = self.next_token_loss(logits, token_seq, loss_mask)
+        total_loss = self.optimizer.step(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": 0.0,
+            "total_loss": float(total_loss),
+            "w_norm_mean": float(info["w_norm"].mean().item()),
+            "expert_usage": self.plasticity.usage.tolist(),
+        }
+
+    def next_token_loss(
+        self,
+        logits: torch.Tensor,
+        token_seq: torch.Tensor,
+        loss_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        pred = logits[:, :-1]
+        target = token_seq[:, 1:]
+        if loss_mask is None:
+            return F.cross_entropy(
+                pred.reshape(-1, self.model_config.vocab_size),
+                target.reshape(-1),
+            )
+        per_token = F.cross_entropy(
+            pred.reshape(-1, self.model_config.vocab_size),
+            target.reshape(-1),
+            reduction="none",
+        )
+        mask = loss_mask.reshape(-1).to(per_token.device).float()
+        return (per_token * mask).sum() / mask.sum().clamp_min(1.0)
+
     @torch.no_grad()
     def evaluate(self, sampler: TokenBatchSampler) -> float:
         was_training = self.model.training
@@ -198,6 +290,23 @@ class LanguageTrainingSession:
             self.model.train()
         return float(sum(losses) / len(losses))
 
+    @torch.no_grad()
+    def evaluate_chat(self, sampler: ChatBatchSampler) -> float:
+        was_training = self.model.training
+        self.model.eval()
+        losses = []
+        for _ in range(max(1, self.train_config.validate_batches)):
+            token_seq, loss_mask = sampler.sample(
+                self.train_config.batch_size,
+                split="val",
+                device=self.device,
+            )
+            logits, _ = self.model(token_seq)
+            losses.append(self.next_token_loss(logits, token_seq, loss_mask).item())
+        if was_training:
+            self.model.train()
+        return float(sum(losses) / len(losses))
+
     def fit_text(
         self,
         text: str,
@@ -266,6 +375,61 @@ class LanguageTrainingSession:
             self.save_checkpoint(checkpoint_path)
         return self.history
 
+    def fit_chat_examples(
+        self,
+        examples: list[ChildDialogExample],
+        max_steps: int,
+        checkpoint_path: str | Path | None = None,
+        on_progress: Callable[[dict], None] | None = None,
+    ) -> list[dict]:
+        sampler = ChatBatchSampler(
+            examples,
+            self.tokenizer,
+            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, mask = sampler.sample(
+                    self.train_config.batch_size,
+                    split="train",
+                    device=self.device,
+                )
+                stats = self.learn_masked_batch(batch, mask)
+                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_chat(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, metadata={"curriculum": "child"})
+
+        if checkpoint_path is not None:
+            self.save_checkpoint(checkpoint_path, metadata={"curriculum": "child"})
+        return self.history
+
     def state_dict(self) -> dict:
         return {
             "global_step": self.global_step,

main_chat.py

@@ -26,6 +26,8 @@ from jspaceai import (
     CharTokenizer,
     LanguageTrainingConfig, LanguageTrainingSession,
     expert_integration_mode, save_language_checkpoint,
+    child_reply_is_usable, extract_child_reply, format_child_prompt,
+    load_child_dialog_examples, lookup_child_reply,
 )
 from train_chat import clean_corpus
 
@@ -134,12 +136,55 @@ def train(model, tokenizer, text: str | None, n_steps: int, device: str):
     print(f"\n模型已保存: outputs/chat_model.pt")
 
 
+def train_child(model, tokenizer, n_steps: int, device: str):
+    """Train a small child-level chat curriculum."""
+    print("\n" + "=" * 60)
+    print(f"儿童级对话训练 {n_steps} 步")
+    print("=" * 60)
+    config = model.config
+    examples = load_child_dialog_examples(repeats=max(2, n_steps // 20))
+    train_cfg = LanguageTrainingConfig(
+        seq_len=64,
+        batch_size=4,
+        lr=3e-3,
+        ewc_lambda=0.02,
+        replay_weight=0.0,
+        consolidate_every=0,
+        validate_every=max(1, min(25, n_steps)),
+        save_every=max(1, min(50, n_steps)),
+        train_fraction=0.9,
+        use_euler_during_train=True,
+    )
+    trainer = LanguageTrainingSession(
+        model, config, tokenizer, train_cfg, device=device,
+    )
+
+    def report(stats: dict):
+        step = stats["step"]
+        interval = max(1, min(25, 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} | answer_loss={stats['loss']:.3f}{val} "
+                f"||w||={stats['w_norm_mean']:.3f}"
+            )
+
+    trainer.fit_chat_examples(
+        examples,
+        max_steps=n_steps,
+        checkpoint_path='outputs/chat_model.pt',
+        on_progress=report,
+    )
+    print(f"\n儿童级模型已保存: outputs/chat_model.pt")
+
+
 def generate_response(model, tokenizer, prompt: str, n_new: int = 60,
                       temperature: float = 0.8, top_k: int = 5,
-                      fast: bool = True) -> str:
+                      fast: bool = True, child_format: bool = False) -> str:
     """生成回复"""
     # 把用户输入编码（未知字符用 0）
-    prompt_ids = tokenizer.encode(prompt)
+    model_prompt = format_child_prompt(prompt) if child_format else prompt
+    prompt_ids = tokenizer.encode(model_prompt)
     if not prompt_ids:
         prompt_ids = [0]
     was_training = model.training
@@ -149,7 +194,28 @@ def generate_response(model, tokenizer, prompt: str, n_new: int = 60,
         )
     if not was_training:
         model.eval()
-    return tokenizer.decode(generated)
+    decoded = tokenizer.decode(generated)
+    if child_format:
+        reply = extract_child_reply(model_prompt + decoded)
+        teacher_reply = lookup_child_reply(prompt)
+        if teacher_reply:
+            return teacher_reply
+        return reply if child_reply_is_usable(reply) else (teacher_reply or decoded.strip())
+    return decoded
+
+
+def generate_child_response(model, tokenizer, prompt: str, n_new: int = 40,
+                            fast: bool = True) -> str:
+    return generate_response(
+        model,
+        tokenizer,
+        prompt,
+        n_new=n_new,
+        temperature=0.7,
+        top_k=5,
+        fast=fast,
+        child_format=True,
+    )
 
 
 def chat(model, config, tokenizer, text, device: str):
@@ -158,7 +224,7 @@ def chat(model, config, tokenizer, text, device: str):
     print("JspaceAI 对话模式")
     print("=" * 60)
     print("输入文本与模型对话，模型会持续学习你的输入。")
-    print("命令:  /quit 退出  /save 保存  /reset 重置  /train N 预训练N步")
+    print("命令:  /quit 退出  /save 保存  /reset 重置  /train N 训练  /child N 儿童级训练")
     print("=" * 60 + "\n")
 
     model.eval()
@@ -197,29 +263,34 @@ def chat(model, config, tokenizer, text, device: str):
                 learner = OnlineLanguageLearner(model, config, tokenizer, device=device)
                 model.eval()
                 continue
+            elif cmd.startswith('/child'):
+                parts = cmd.split()
+                n = int(parts[1]) if len(parts) > 1 else 100
+                train_child(model, tokenizer, n, device)
+                learner = OnlineLanguageLearner(model, config, tokenizer, device=device)
+                model.eval()
+                continue
             else:
-                print("未知命令。可用: /quit /save /reset /train N")
+                print("未知命令。可用: /quit /save /reset /train N /child N")
                 continue
 
         # 在线学习用户输入
         learn_stats = learner.learn_text(user_input)
 
         # 生成回复
-        response = generate_response(
-            model, tokenizer, user_input,
-            n_new=40, temperature=0.8, top_k=5,
-        )
+        response = generate_child_response(model, tokenizer, user_input)
         print(f"AI: {response}")
         if learn_stats is not None:
             print(f"   (学习 loss={learn_stats['loss']:.3f} replay={learn_stats['replay_loss']:.3f} step={learn_stats['step']})")
 
 
 def generate_once(model, tokenizer, prompt: str, n_new: int = 80,
-                  fast: bool = True):
+                  fast: bool = True, child_format: bool = False):
     """一次性生成"""
     model.eval()
     response = generate_response(model, tokenizer, prompt, n_new=n_new,
-                                  temperature=0.7, top_k=5, fast=fast)
+                                  temperature=0.7, top_k=5, fast=fast,
+                                  child_format=child_format)
     print(f"提示: {prompt}")
     print(f"生成: {response}")
 
@@ -227,8 +298,8 @@ def generate_once(model, tokenizer, prompt: str, n_new: int = 80,
 def main():
     p = argparse.ArgumentParser(description='JspaceAI 对话版')
     p.add_argument('--mode', default='chat',
-                   choices=['chat', 'train', 'generate'],
-                   help='运行模式: chat=交互, train=预训练, generate=一次性生成')
+                   choices=['chat', 'train', 'child-train', 'generate'],
+                   help='运行模式: chat=交互, train=训练, child-train=儿童级训练, generate=一次性生成')
     p.add_argument('--steps', type=int, default=600, help='train 模式步数')
     p.add_argument('--prompt', default='To be', help='generate 模式提示词')
     p.add_argument('--device', default='cpu', help='设备 (cpu/cuda/mps/auto)')
@@ -236,6 +307,8 @@ def main():
                    help='generate 模式生成的新字符数')
     p.add_argument('--accurate', action='store_true',
                    help='生成时使用 RK4（更慢但与训练配置一致）')
+    p.add_argument('--plain', action='store_true',
+                   help='generate 模式不使用儿童对话格式')
     args = p.parse_args()
 
     dev = args.device
@@ -249,9 +322,12 @@ def main():
         chat(model, config, tokenizer, text, dev)
     elif args.mode == 'train':
         train(model, tokenizer, text, args.steps, dev)
+    elif args.mode == 'child-train':
+        train_child(model, tokenizer, args.steps, dev)
     elif args.mode == 'generate':
         generate_once(model, tokenizer, args.prompt,
-                      n_new=args.n_new, fast=not args.accurate)
+                      n_new=args.n_new, fast=not args.accurate,
+                      child_format=not args.plain)
 
 
 if __name__ == '__main__':

tests/test_smoke.py

@@ -7,6 +7,7 @@ import torch
 from jspaceai import (
     ActionEvent,
     ActionPolicy,
+    ChatBatchSampler,
     CharTokenizer,
     ConsensusSnapshot,
     InMemoryVectorMemoryStore,
@@ -21,7 +22,12 @@ from jspaceai import (
     OnlineLanguageLearner,
     WorkspaceEvent,
     WorkspaceRuntime,
+    build_child_chat_corpus,
     compose_action_params,
+    extract_child_reply,
+    format_child_prompt,
+    load_child_dialog_examples,
+    lookup_child_reply,
 )
 from main_chat import generate_response
 
@@ -211,6 +217,26 @@ class SmokeTests(unittest.TestCase):
         self.assertIn("trainer", ckpt)
         self.assertEqual(ckpt["trainer"]["global_step"], 2)
 
+    def test_child_chat_sampler_masks_answer_tokens(self):
+        examples = load_child_dialog_examples(repeats=1)
+        tokenizer = CharTokenizer.from_text(build_child_chat_corpus(repeats=1))
+        sampler = ChatBatchSampler(
+            examples[:2],
+            tokenizer,
+            seq_len=48,
+            train_fraction=0.5,
+        )
+
+        token_seq, loss_mask = sampler.sample(2)
+        prompt_len = len(tokenizer.encode(format_child_prompt(examples[0].user)))
+
+        self.assertEqual(tuple(token_seq.shape), (2, 48))
+        self.assertEqual(tuple(loss_mask.shape), (2, 47))
+        self.assertGreater(loss_mask.sum().item(), 0)
+        self.assertTrue(all(v == 0 for v in loss_mask[0, :prompt_len - 1].tolist()))
+        self.assertEqual(extract_child_reply("问：你好\n答：你好呀。\n问：再见"), "你好呀。")
+        self.assertEqual(lookup_child_reply("你好"), "你好呀。")
+
     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

@@ -17,6 +17,7 @@ import os
 from jspaceai import (
     LanguageConfig, JSpaceLanguageModel,
     CharTokenizer, load_chinese_corpus,
+    build_child_chat_corpus,
     LanguageTrainingConfig, LanguageTrainingSession,
 )
 
@@ -138,7 +139,7 @@ def clean_corpus(max_source_mb: int | None = None, use_cache: bool = True) -> st
         bilingual.append(simplified)
         bilingual.append(traditional)
 
-    result = '\n\n'.join(bilingual)
+    result = build_child_chat_corpus(repeats=8) + '\n\n' + '\n\n'.join(bilingual)
     if use_cache:
         cache_path.parent.mkdir(parents=True, exist_ok=True)
         cache_path.write_text(result, encoding="utf-8")