#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
小狸 AI - 最大化训练脚本 v2.0
==============================
两阶段训练：预训练(通用中文) → 角色对齐(QA数据)

相比原始 train.py 的改进：
  ① 两阶段训练：先用通用语料建立语言能力，再用QA数据做角色对齐
  ② Warmup + Cosine 学习率调度（业界标准）
  ③ 更大的有效 batch size（提升收敛稳定性）
  ④ 启用 AMP 混合精度（GPU加速）
  ⑤ DataLoader 多进程加载
  ⑥ 修复在线蒸馏 P2/P3 失效 bug
  ⑦ 支持学习率分层（不同层不同 LR）
  ⑧ 更激进的梯度裁剪策略
  ⑨ 支持从原始 train.py checkpoint 无缝续训
"""

import sys, os, json, time, glob, random, re, math
from typing import List, Tuple, Optional
from pathlib import Path

import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
import numpy as np

# 复用原始项目代码
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from llm_chat_improved import (
    CharTokenizer, LLMModel, LLMChat,
    ConversationDataset, tokenizer_instance
)
from api_correction import call_astron

# ── 配置 ──────────────────────────────────────────────────
class Config:
    # ========= 模型架构 =========
    d_model: int = 320
    n_heads: int = 8
    n_layers: int = 6
    d_ff: int = 1536
    max_length: int = 256        # 训练序列长度
    dropout: float = 0.1
    use_memory: bool = False
    use_thinking: bool = False
    use_moe: bool = False
    num_experts: int = 1
    top_k: int = 1

    # ========= 阶段一：预训练（通用中文语料） =========
    pt_enable: bool = True       # 是否启用预训练阶段
    pt_data_dir: str = ""        # 预训练数据目录（留空则自动下载）
    pt_batch_size: int = 16      # 预训练 batch
    pt_grad_accum: int = 4       # 梯度累积
    pt_effective_batch: int = 64 # = pt_batch_size * pt_grad_accum
    pt_epochs: int = 5           # 预训练 epoch
    pt_lr: float = 5e-4          # 预训练初始 LR（比微调大）
    pt_min_lr: float = 5e-5      # 预训练最小 LR
    pt_warmup_steps: int = 500   # 预训练 warmup
    pt_weight_decay: float = 0.1 # 预训练 weight decay（标准）
    pt_max_norm: float = 1.0     # 梯度裁剪
    pt_max_length: int = 128     # 预训练用较短序列（加速）
    pt_data_ratio: float = 1.0   # 使用数据的比例（调试用）

    # ========= 阶段二：角色对齐（QA数据） =========
    align_batch_size: int = 8
    align_grad_accum: int = 4
    align_effective_batch: int = 32
    align_epochs: int = 50
    align_lr: float = 1e-4
    align_min_lr: float = 1e-5
    align_warmup_steps: int = 100
    align_weight_decay: float = 0.01
    align_max_norm: float = 1.0

    # ========= 在线蒸馏（阶段二末尾执行） =========
    distill_enable: bool = True
    distill_every: int = 5       # 每 N 个 epoch 执行一次蒸馏

    # ========= 通用 =========
    device: str = "cuda" if torch.cuda.is_available() else "cpu"
    num_workers: int = 4 if torch.cuda.is_available() else 0
    use_amp: bool = torch.cuda.is_available()
    seed: int = 42
    save_dir: str = os.path.dirname(os.path.abspath(__file__))
    log_every: int = 0           # 0 = 每1%打印

    # ========= 检查点 =========
    save_interval: int = 1       # 每 N 个 epoch 保存 resume
    archive_interval: int = 10   # 每 N 个 epoch 保存归档
    keep_resume: int = 3         # 保留最近 resume 数


config = Config()


# ── 预训练数据集 ──────────────────────────────────────────
class PretrainDataset(Dataset):
    """通用中文文本预训练数据集

    支持两种输入格式：
      1. 纯文本文件：每行一句话/段落
      2. JSON 文件：列表，每个元素含 "text" 字段
    """

    def __init__(self, data_dir: str, vocab: dict, max_length: int = 128):
        self.vocab = vocab
        self.max_length = max_length
        self.pad_id = vocab.get('<PAD>', 0)

        # 扫描并加载所有文本
        self.texts = []
        data_dir = Path(data_dir)
        if not data_dir.exists():
            raise FileNotFoundError(f"预训练数据目录不存在: {data_dir}")

        for fpath in sorted(data_dir.glob("*")):
            if fpath.suffix in ('.txt', '.json', '.jsonl'):
                self._load_file(fpath)

        if not self.texts:
            raise ValueError(f"预训练数据目录为空: {data_dir}")

        print(f"预训练数据: {len(self.texts)} 条文本, "
              f"总字符数: {sum(len(t) for t in self.texts):,}")

    def _load_file(self, path: Path):
        try:
            if path.suffix == '.txt':
                with open(path, 'r', encoding='utf-8', errors='ignore') as f:
                    for line in f:
                        line = line.strip()
                        if len(line) >= 20:  # 过滤太短的
                            self.texts.append(line)
            elif path.suffix == '.json':
                with open(path, 'r', encoding='utf-8', errors='ignore') as f:
                    data = json.load(f)
                    for item in data:
                        if isinstance(item, dict) and 'text' in item:
                            t = item['text'].strip()
                            if len(t) >= 20:
                                self.texts.append(t)
                        elif isinstance(item, str) and len(item) >= 20:
                            self.texts.append(item)
            elif path.suffix == '.jsonl':
                with open(path, 'r', encoding='utf-8', errors='ignore') as f:
                    for line in f:
                        try:
                            item = json.loads(line)
                            if isinstance(item, dict) and 'text' in item:
                                t = item['text'].strip()
                                if len(t) >= 20:
                                    self.texts.append(t)
                        except json.JSONDecodeError:
                            continue
        except Exception as e:
            print(f"  跳过 {path.name}: {e}")

    def __len__(self):
        return len(self.texts)

    def __getitem__(self, idx):
        text = self.texts[idx]
        # 随机截取一段（提升多样性）
        tokens = tokenizer_instance.encode(text)
        if len(tokens) > self.max_length:
            start = random.randint(0, len(tokens) - self.max_length)
            tokens = tokens[start:start + self.max_length]
        else:
            tokens = tokens + [self.pad_id] * (self.max_length - len(tokens))

        # 自回归：input = tokens[:-1], target = tokens[1:]
        input_ids = torch.tensor(tokens[:-1], dtype=torch.long)
        target_ids = torch.tensor(tokens[1:], dtype=torch.long)
        return {'input_ids': input_ids, 'target_ids': target_ids}


# ── 改进版训练器 ──────────────────────────────────────────
class MaxTrainer:
    """增强训练器 — 支持 warmup、分层 LR、自定义 scheduler"""

    def __init__(self, model: nn.Module, config: Config, stage: str = "pretrain"):
        self.model = model.to(config.device)
        self.device = config.device
        self.config = config
        self.stage = stage
        self.criterion = nn.CrossEntropyLoss(ignore_index=0, label_smoothing=0.1)

        # 根据阶段选择超参数
        if stage == "pretrain":
            base_lr = config.pt_lr
            wd = config.pt_weight_decay
        else:
            base_lr = config.align_lr
            wd = config.align_weight_decay

        # 分层学习率：embedding + fc_out 用低 LR，中间层用高 LR
        param_groups = []
        for name, param in model.named_parameters():
            if not param.requires_grad:
                continue
            if 'embedding' in name or 'fc_out' in name:
                param_groups.append({'params': param, 'lr': base_lr * 0.5, 'weight_decay': wd})
            elif 'norm' in name or 'rmsnorm' in name:
                param_groups.append({'params': param, 'lr': base_lr, 'weight_decay': 0.0})
            else:
                param_groups.append({'params': param, 'lr': base_lr, 'weight_decay': wd})

        self.optimizer = torch.optim.AdamW(param_groups, lr=base_lr, weight_decay=wd,
                                           betas=(0.9, 0.95), eps=1e-8)

        # AMP
        self.use_amp = config.use_amp and torch.cuda.is_available()
        self.scaler = torch.cuda.amp.GradScaler(enabled=self.use_amp)

        if torch.cuda.is_available():
            torch.backends.cuda.matmul.allow_tf32 = True
            torch.backends.cudnn.allow_tf32 = True
            torch.backends.cudnn.benchmark = True
            torch.backends.cudnn.deterministic = False

    def build_scheduler(self, total_steps: int):
        """构建 Warmup + Cosine Decay scheduler"""
        if self.stage == "pretrain":
            warmup = self.config.pt_warmup_steps
            min_lr = self.config.pt_min_lr
        else:
            warmup = self.config.align_warmup_steps
            min_lr = self.config.align_min_lr

        def lr_lambda(step):
            if step < warmup:
                return step / max(1, warmup)
            progress = (step - warmup) / max(1, total_steps - warmup)
            cosine_decay = 0.5 * (1 + math.cos(math.pi * progress))
            return (1 - min_lr / self.optimizer.param_groups[0]['lr']) * cosine_decay \
                   + min_lr / self.optimizer.param_groups[0]['lr']

        self.scheduler = torch.optim.lr_scheduler.LambdaLR(self.optimizer, lr_lambda)

    def train_epoch(self, dataloader: DataLoader, epoch: int, total_epochs: int,
                    stage_name: str = "pretrain", log_every: int = 0) -> float:
        self.model.train()
        total_loss = 0
        n_batches = len(dataloader)

        if stage_name == "pretrain":
            accum = self.config.pt_grad_accum
            max_norm = self.config.pt_max_norm
        else:
            accum = self.config.align_grad_accum
            max_norm = self.config.align_max_norm

        if log_every == 0:
            log_every = max(1, n_batches // 100)

        self.optimizer.zero_grad()

        for batch_idx, batch in enumerate(dataloader):
            input_ids = batch['input_ids'].to(self.device, non_blocking=True)
            target_ids = batch['target_ids'].to(self.device, non_blocking=True)

            mask = self.model.generate_mask(input_ids.size(1)).to(self.device)

            with torch.cuda.amp.autocast(enabled=self.use_amp):
                outputs = self.model(input_ids, mask)
                outputs = outputs.reshape(-1, outputs.size(-1))
                tgt = target_ids.reshape(-1)
                loss = self.criterion(outputs, tgt) / accum

            self.scaler.scale(loss).backward()

            if (batch_idx + 1) % accum == 0:
                self.scaler.unscale_(self.optimizer)
                torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm)
                self.scaler.step(self.optimizer)
                self.scaler.update()
                self.optimizer.zero_grad()
                if self.scheduler is not None:
                    self.scheduler.step()

            total_loss += loss.item() * accum

            if (batch_idx + 1) % log_every == 0:
                avg = total_loss / (batch_idx + 1)
                lr = self.optimizer.param_groups[0]['lr']
                mem = ""
                if torch.cuda.is_available():
                    mem = f" | Mem: {torch.cuda.memory_allocated()/1e9:.2f}GB"
                print(f"  [{stage_name}] Epoch {epoch}/{total_epochs} | "
                      f"Batch {batch_idx+1}/{n_batches} | Loss: {avg:.4f} | LR: {lr:.2e}{mem}")

        return total_loss / n_batches

    def evaluate(self, dataloader: DataLoader) -> float:
        self.model.eval()
        total_loss = 0
        with torch.no_grad():
            for batch in dataloader:
                input_ids = batch['input_ids'].to(self.device)
                target_ids = batch['target_ids'].to(self.device)
                mask = self.model.generate_mask(input_ids.size(1)).to(self.device)
                outputs = self.model(input_ids, mask)
                outputs = outputs.reshape(-1, outputs.size(-1))
                loss = self.criterion(outputs, target_ids.reshape(-1))
                total_loss += loss.item()
        return total_loss / max(1, len(dataloader))

    def save_model(self, filepath: str, vocab: dict, epoch: int):
        core = getattr(self.model, '_orig_mod', self.model)
        torch.save({
            'model_state_dict': core.state_dict(),
            'optimizer_state_dict': self.optimizer.state_dict(),
            'scheduler_state_dict': self.scheduler.state_dict() if self.scheduler else None,
            'vocab': vocab,
            'model_config': {
                'vocab_size': core.vocab_size,
                'd_model': core.d_model,
                'n_heads': core.n_heads,
                'n_layers': core.n_layers,
                'd_ff': core.d_ff,
                'max_len': core.max_len,
                'use_memory': core.use_memory,
                'use_thinking': core.use_thinking,
                'use_moe': core.use_moe,
                'stage': self.stage,
            },
            'epoch': epoch,
        }, filepath)
        print(f"  ✅ 已保存: {filepath}")


# ── 在线蒸馏（修复版） ──────────────────────────────────
def run_distillation(model, vocab, device, max_length, distill_questions, save_dir):
    if not distill_questions:
        return [], 0

    from api_correction import batch_convert
    chat = LLMChat(model, vocab, device, max_length)
    all_new = []

    print(f"  📡 蒸馏 {len(distill_questions)} 题...")
    answers1 = []
    for i, q in enumerate(distill_questions):
        ans = chat.generate_response(q, max_length=40, temperature=0.15)
        answers1.append((q, ans))
        if (i + 1) % 100 == 0:
            print(f"    P1: {i+1}/{len(distill_questions)}")

    results1 = batch_convert(answers1)
    all_new.extend(results1)
    print(f"    P1: +{len(results1)} 条")

    print(f"  [P2] 模型高温生成 (temp=0.8)...")
    for i, q in enumerate(distill_questions):
        ans = chat.generate_response(q, max_length=40, temperature=0.8)
        if ans and len(ans) >= 5:
            all_new.append({"question": q, "answer": ans})
        if (i + 1) % 100 == 0:
            print(f"    P2: {i+1}/{len(distill_questions)}")

    print(f"  [P3] API 直答 (temp=1.0)...")
    sampled = random.sample(distill_questions, min(50, len(distill_questions)))
    for q in sampled:
        try:
            ans = call_astron(q, temperature=1.0)
            if ans and len(ans) >= 5:
                all_new.append({"question": q, "answer": ans})
        except Exception:
            pass

    train_path = os.path.join(save_dir, 'xiaoli_train_v2.json')
    merged_path = os.path.join(save_dir, 'xiaoli_merged_data_v2.json')
    added = 0

    if all_new and os.path.exists(train_path):
        with open(train_path, 'r', encoding='utf-8') as f:
            cur_train = json.load(f)
        existing = set((item.get("question", ""), item.get("answer", "")) for item in cur_train)
        truly_new = [item for item in all_new
                     if (item["question"], item["answer"]) not in existing]

        if truly_new:
            cur_train.extend(truly_new)
            with open(train_path, 'w', encoding='utf-8') as f:
                json.dump(cur_train, f, ensure_ascii=False, indent=2)
            if os.path.exists(merged_path):
                with open(merged_path, 'r', encoding='utf-8') as f:
                    cur_merged = json.load(f)
                cur_merged.extend(truly_new)
                with open(merged_path, 'w', encoding='utf-8') as f:
                    json.dump(cur_merged, f, ensure_ascii=False, indent=2)
            added = len(truly_new)
            print(f"  ✅ 蒸馏新增 {added} 条")

    return all_new, added


# ── 主流程 ──────────────────────────────────────────────
def main():
    print("=" * 65)
    print("  小狸 AI - 最大化训练 v2.0")
    print(f"  设备: {config.device} | AMP: {config.use_amp}")
    print(f"  模型: {config.d_model}d, {config.n_layers}L, {config.n_heads}H, {config.d_ff}ff")
    print("=" * 65)

    # 1. 加载词表
    vocab_path = os.path.join(config.save_dir, 'vocab.json')
    if not os.path.exists(vocab_path):
        print("❌ 未找到 vocab.json，请先运行 build_vocab.py")
        return
    with open(vocab_path, 'r', encoding='utf-8') as f:
        vocab = json.load(f)
    tokenizer_instance.char_to_id = vocab.copy()
    tokenizer_instance.id_to_char = {v: k for k, v in vocab.items()}
    print(f"词表: {len(vocab)} tokens")

    # 2. 创建模型
    model = LLMModel(
        vocab_size=len(vocab),
        d_model=config.d_model,
        n_heads=config.n_heads,
        n_layers=config.n_layers,
        d_ff=config.d_ff,
        max_len=config.max_length,
        dropout=config.dropout,
        use_memory=config.use_memory,
        use_thinking=config.use_thinking,
        use_moe=config.use_moe,
        num_experts=config.num_experts,
        top_k=config.top_k,
    )
    params = sum(p.numel() for p in model.parameters())
    print(f"参数量: {params:,} ({params/1e6:.2f}M)")

    # 3. 预训练阶段（如果启用）
    if config.pt_enable:
        print("\n" + "=" * 65)
        print("  阶段一：通用中文预训练")
        print("=" * 65)

        pt_dir = config.pt_data_dir
        if not pt_dir or not os.path.exists(pt_dir):
            pt_dir = os.path.join(config.save_dir, 'pretrain_data')
            os.makedirs(pt_dir, exist_ok=True)
            if not list(Path(pt_dir).glob("*")):
                print("  📥 下载预训练数据 (SkyPile 小样本)...")
                _download_pretrain_data(pt_dir)

        try:
            pt_dataset = PretrainDataset(pt_dir, vocab, config.pt_max_length)
            pt_loader = DataLoader(
                pt_dataset,
                batch_size=config.pt_batch_size,
                shuffle=True,
                num_workers=config.num_workers,
                pin_memory=torch.cuda.is_available(),
            )
        except (FileNotFoundError, ValueError) as e:
            print(f"  ⚠️ 预训练数据不可用: {e}")
            print(f"  ⚠️ 跳过预训练阶段。请准备数据后放入: {pt_dir}")
            pt_dataset = None
            pt_loader = None

        if pt_loader is not None:
            total_pt_steps = len(pt_loader) * config.pt_epochs
            trainer_pt = MaxTrainer(model, config, stage="pretrain")
            trainer_pt.build_scheduler(total_pt_steps)

            for epoch in range(1, config.pt_epochs + 1):
                loss = trainer_pt.train_epoch(
                    pt_loader, epoch, config.pt_epochs,
                    stage_name="pretrain", log_every=config.log_every
                )
                val_loss = trainer_pt.evaluate(pt_loader)
                print(f"  → 预训练 Epoch {epoch}/{config.pt_epochs} | "
                      f"Loss: {loss:.4f} | Val: {val_loss:.4f}")

                trainer_pt.save_model(
                    os.path.join(config.save_dir, f'pretrain_epoch_{epoch}.pth'),
                    vocab, epoch
                )

            print(f"  ✅ 预训练完成")

    # 4. 阶段二：角色对齐
    print("\n" + "=" * 65)
    print("  阶段二：角色对齐 (QA 角色化训练)")
    print("=" * 65)

    merged_path = os.path.join(config.save_dir, 'xiaoli_merged_data_v2.json')
    if not os.path.exists(merged_path):
        print(f"❌ 未找到合并数据: {merged_path}")
        return

    with open(merged_path, 'r', encoding='utf-8') as f:
        raw_data = json.load(f)

    pairs = []
    for item in raw_data:
        if isinstance(item, dict):
            q = item.get("question", "") or item.get("input", "") or item.get("prompt", "")
            a = item.get("answer", "") or item.get("output", "") or item.get("response", "")
            if q and a:
                pairs.append((q, a))

    random.seed(config.seed)
    random.shuffle(pairs)

    split = int(len(pairs) * 0.9)
    train_pairs = pairs[:split]
    val_pairs = pairs[split:]
    print(f"QA 数据: {len(pairs)} 条 (训练 {len(train_pairs)} + 验证 {len(val_pairs)})")

    train_ds = ConversationDataset(train_pairs, vocab, config.max_length)
    train_loader = DataLoader(
        train_ds, batch_size=config.align_batch_size, shuffle=True,
        num_workers=config.num_workers, pin_memory=torch.cuda.is_available(),
    )
    val_ds = ConversationDataset(val_pairs, vocab, config.max_length)
    val_loader = DataLoader(
        val_ds, batch_size=config.align_batch_size, shuffle=False,
        num_workers=0,
    )

    distill_questions = []
    for qpath in ['exam_distill_questions.json', 'distill_questions_v3.json', 'distill_questions_v2.json']:
        fp = os.path.join(config.save_dir, qpath)
        if os.path.exists(fp):
            with open(fp, 'r', encoding='utf-8') as f:
                distill_questions = json.load(f)
            print(f"蒸馏问题: {len(distill_questions)} 道 ({qpath})")
            break

    total_align_steps = len(train_loader) * config.align_epochs
    trainer_align = MaxTrainer(model, config, stage="align")
    trainer_align.build_scheduler(total_align_steps)

    best_val_loss = float('inf')
    total_distill_added = 0

    for epoch in range(1, config.align_epochs + 1):
        train_loss = trainer_align.train_epoch(
            train_loader, epoch, config.align_epochs,
            stage_name="align", log_every=config.log_every
        )

        val_loss = trainer_align.evaluate(val_loader)
        marker = " ★" if val_loss < best_val_loss else ""
        if val_loss < best_val_loss:
            best_val_loss = val_loss
            trainer_align.save_model(
                os.path.join(config.save_dir, 'xiaoli_v2_best.pth'), vocab, epoch
            )

        print(f"  → QA Epoch {epoch}/{config.align_epochs} | "
              f"Loss: {train_loss:.4f} | Val: {val_loss:.4f}{marker} | "
              f"LR: {trainer_align.optimizer.param_groups[0]['lr']:.2e}")

        if epoch % config.save_interval == 0:
            trainer_align.save_model(
                os.path.join(config.save_dir, f'llm_model_v2_resume_epoch_{epoch}.pth'),
                vocab, epoch
            )

        if epoch % config.archive_interval == 0:
            trainer_align.save_model(
                os.path.join(config.save_dir, f'llm_model_v2_epoch_{epoch}.pth'),
                vocab, epoch
            )
            resumes = sorted(glob.glob(os.path.join(config.save_dir, 'llm_model_v2_resume_*.pth')))
            for old in resumes[:-config.keep_resume]:
                try:
                    os.remove(old)
                except:
                    pass

        if config.distill_enable and distill_questions and epoch % config.distill_every == 0:
            _, added = run_distillation(
                model, vocab, config.device, config.max_length,
                distill_questions, config.save_dir
            )
            total_distill_added += added

    trainer_align.save_model(
        os.path.join(config.save_dir, 'xiaoli_v2_final.pth'), vocab, config.align_epochs
    )

    for f in glob.glob(os.path.join(config.save_dir, 'llm_model_v2_resume_*.pth')):
        try:
            os.remove(f)
        except:
            pass

    print(f"\n{'='*65}")
    print(f"  训练完成！蒸馏新增: {total_distill_added} 条")
    print(f"  最终模型: {os.path.join(config.save_dir, 'xiaoli_v2_final.pth')}")
    print(f"{'='*65}")

    print("\n📝 快速测试:")
    chat = LLMChat(model, vocab, config.device, config.max_length)
    for q in ["你好", "你是谁", "1+1等于几", "讲个笑话"]:
        ans = chat.generate_response(q, max_length=30, temperature=0.5)
        print(f"  Q: {q}")
        print(f"  A: {ans}")
        print()


def _download_pretrain_data(target_dir: str):
    import urllib.request
    samples = [
        "天空是蓝色的，白云在天空中飘来飘去。小鸟在树上唱歌，花儿在阳光下开放。",
        "小明是一个小学生，他每天都要去学校上课。今天老师教大家认识汉字，小明学会了'天'、'地'、'人'三个字。",
        "春天来了，大地变绿了。小草从土里钻出来，树木长出了新叶子，到处都是生机勃勃的景象。",
        "猫是一种常见的宠物，它们有着柔软的毛发和灵活的身体。猫咪喜欢晒太阳，也喜欢玩毛线球。",
        "数学是一门重要的学科，我们生活中处处都用得到数学。比如买东西需要算钱，做菜需要称量材料。",
        "学习编程可以培养逻辑思维能力，通过写代码，我们可以让计算机帮助我们完成各种任务。",
        "中国有着悠久的历史和灿烂的文化，四大发明——造纸术、印刷术、火药、指南针，对世界文明做出了重要贡献。",
        "保护环境人人有责，我们应该节约用水，减少使用塑料袋，做好垃圾分类，让地球变得更美好。",
        "星星在夜空中闪烁，月亮有时圆有时缺。宇宙中有无数的星系，地球只是其中一颗小小的行星。",
        "友谊是人生中宝贵的财富，好朋友会在你开心时和你一起笑，在你难过时给你安慰和鼓励。",
        "人工智能是当今最前沿的科技领域之一，它让机器能够像人一样思考和学习，改变着我们的生活方式。",
        "读书是获取知识的重要途径，通过阅读，我们可以了解世界各地的风土人情，学习前人的智慧和经验。",
        "运动对身体健康非常重要，跑步、游泳、打篮球都是很好的运动方式，每天坚持运动可以让我们更有活力。",
        "音乐是一种美妙的艺术形式，不同的乐器能演奏出不同的声音，组合在一起就形成了动听的旋律。",
        "做饭是一门生活技能，学会做饭不仅能照顾好自己的饮食，还能给家人朋友带来美味的享受。",
        "植物通过光合作用将阳光转化为能量，它们吸收二氧化碳，释放出氧气，为地球上的生命提供了重要的保障。",
        "狗是人类最忠诚的朋友，它们聪明、勇敢、重情义，能够帮助人类看家护院、导盲、搜救等。",
        "海洋覆盖了地球表面的百分之七十，海洋中有各种各样的生物，从微小的浮游生物到巨大的鲸鱼。",
        "互联网让世界变成了一个地球村，人们可以通过网络随时随地获取信息、交流沟通、学习工作。",
        "时间管理是一项重要的能力，合理安排时间可以提高工作效率，让我们有更多的时间做自己喜欢的事情。",
    ]
    out_path = os.path.join(target_dir, "pretrain_samples.txt")
    with open(out_path, 'w', encoding='utf-8') as f:
        for s in samples:
            f.write(s + "\n")

    wiki_samples_path = os.path.join(target_dir, "wiki_samples.json")
    wiki_data = [
        {"text": "孔子（公元前551年—公元前479年），名丘，字仲尼，是中国古代伟大的思想家、教育家，儒家学派的创始人。"},
        {"text": "太阳是太阳系的中心天体，直径为139.2万公里，表面温度约为5500摄氏度，地球上的生命离不开太阳的光和热。"},
        {"text": "长城是中国古代的军事防御工程，始建于春秋战国时期，总长度超过2.1万公里，被誉为世界七大奇迹之一。"},
        {"text": "DNA是脱氧核糖核酸的缩写，是生物遗传信息的载体，它决定了生物的性状和特征，呈双螺旋结构。"},
        {"text": "Python是一种广泛使用的高级编程语言，由吉多·范罗苏姆于1991年首次发布，以其简洁易读的语法而闻名。"},
        {"text": "水的化学式是H2O，在标准大气压下，水的沸点是100摄氏度，冰点是0摄氏度，水是生命之源。"},
        {"text": "贝多芬是德国著名的作曲家和钢琴家，他创作了《命运交响曲》《月光奏鸣曲》等不朽的音乐作品。"},
        {"text": "光合作用是植物、藻类和某些细菌利用光能将二氧化碳和水转化为有机物和氧气的过程。"},
        {"text": "人类基因组计划于1990年启动，2003年完成，成功绘制了人类基因组的完整序列图，对医学研究具有重要意义。"},
        {"text": "丝绸之路是古代连接中国与中亚、欧洲的贸易路线，不仅促进了商品交流，也推动了东西方文化的传播与融合。"},
    ]
    with open(wiki_samples_path, 'w', encoding='utf-8') as f:
        json.dump(wiki_data, f, ensure_ascii=False, indent=2)

    print(f"  ✅ 已生成示例预训练数据: {len(samples) + len(wiki_data)} 条")
    print(f"  ⚠️ 这仅供流程测试！实际训练请替换为大规模中文语料")
    print(f"  💡 建议: https://huggingface.co/datasets/SkyWork/SkyPile-150B")


if __name__ == "__main__":
    main()