221 lines
7.6 KiB
Python
221 lines
7.6 KiB
Python
#!/usr/bin/env python3
|
||
"""训练对话模型——清洗语料 + 可恢复分阶段训练
|
||
|
||
策略：
|
||
1. 递归读取 corpus/ 下的本地语料
|
||
2. 清理 markdown/HTML 噪声，过滤过短或过长段落
|
||
3. 同时投喂简体和繁体版本
|
||
4. 统一使用 LanguageTrainingSession 训练、验证和保存
|
||
5. 分阶段降低学习率
|
||
"""
|
||
import torch
|
||
from pathlib import Path
|
||
import time
|
||
import re
|
||
import os
|
||
|
||
from jspaceai import (
|
||
LanguageConfig, JSpaceLanguageModel,
|
||
CharTokenizer, load_chinese_corpus,
|
||
LanguageTrainingConfig, LanguageTrainingSession,
|
||
)
|
||
|
||
|
||
def get_config(vocab_size: int) -> LanguageConfig:
|
||
return LanguageConfig(
|
||
vocab_size=vocab_size, embed_dim=64, input_dim=32,
|
||
workspace_dim=128, expert_dim=64, num_experts=12, num_wells=8,
|
||
ode_steps=4, dt=0.1, tau_w=0.6, jacobian_sparsity=32, noise_std=0.001,
|
||
use_rk4=True, use_layer_norm=True,
|
||
)
|
||
|
||
|
||
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/ 目录下的本地语料（默认有读取预算）
|
||
2. 去掉 markdown/HTML 标记（链接、表格、标题符号等），只保留纯文本
|
||
3. 对每段文本同时生成简体版和繁体版，都喂给模型
|
||
4. 加上内嵌唐诗宋词论语（简体连续文本）
|
||
"""
|
||
from opencc import OpenCC
|
||
|
||
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'):
|
||
para = para.strip()
|
||
if para:
|
||
raw_paragraphs.append(para)
|
||
|
||
# 2. LCCC 对话语料（最高优先级，真实对话）
|
||
lccc_file = corpus_dir / 'lccc_dialogues.txt'
|
||
if lccc_file.exists():
|
||
lccc_text = lccc_file.read_text(encoding='utf-8')
|
||
for para in lccc_text.split('\n\n'):
|
||
para = para.strip()
|
||
if para:
|
||
raw_paragraphs.append(para)
|
||
|
||
# 3. 递归读取 corpus/ 下其他所有文件（教材+百科）
|
||
if corpus_dir.exists():
|
||
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) # 标题
|
||
content = re.sub(r'^\|.*\|$', '', content, flags=re.M) # 表格行
|
||
content = re.sub(r'^---+$', '', content, flags=re.M) # 分隔线
|
||
content = re.sub(r'```[^`]*```', '', content, flags=re.S) # 代码块
|
||
content = re.sub(r'\*\*([^*]+)\*\*', r'\1', content) # 粗体
|
||
content = re.sub(r'\*([^*]+)\*', r'\1', content) # 斜体
|
||
content = re.sub(r'^- ', '', content, flags=re.M) # 列表
|
||
# 按空行分段
|
||
for para in content.split('\n\n'):
|
||
para = para.strip()
|
||
if para:
|
||
raw_paragraphs.append(para)
|
||
|
||
# 清洗过滤
|
||
cleaned = []
|
||
for para in raw_paragraphs:
|
||
para = re.sub(r'[ \t]+', ' ', para)
|
||
para = re.sub(r'\n{3,}', '\n\n', para)
|
||
para = para.strip()
|
||
if len(para) < 10 or len(para) > 500:
|
||
continue
|
||
# 至少要有一些中文字符
|
||
zh_chars = sum(1 for c in para if len(c) == 1 and ord(c) > 0x4e00)
|
||
if zh_chars < 3:
|
||
continue
|
||
cleaned.append(para)
|
||
|
||
# 繁简双版本投喂
|
||
bilingual = []
|
||
for para in cleaned:
|
||
simplified = cc_t2s.convert(para)
|
||
traditional = cc_s2t.convert(simplified)
|
||
bilingual.append(simplified)
|
||
bilingual.append(traditional)
|
||
|
||
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():
|
||
text = clean_corpus()
|
||
print(f"清洗后语料: {len(text)} 字符")
|
||
|
||
tok = CharTokenizer.from_text(text, max_vocab=3000) # 覆盖98.9%文本
|
||
cfg = get_config(tok.vocab_size)
|
||
model = JSpaceLanguageModel(cfg)
|
||
print(f"vocab={tok.vocab_size}, params={sum(p.numel() for p in model.parameters()):,}")
|
||
|
||
all_tokens = tok.encode(text)
|
||
|
||
# 加载已有模型
|
||
mp = Path('outputs/chat_model.pt')
|
||
if mp.exists():
|
||
try:
|
||
ckpt = torch.load(mp, map_location='cpu', weights_only=False)
|
||
if ckpt['config'].vocab_size == cfg.vocab_size:
|
||
model.load_state_dict(ckpt['model'])
|
||
print(f"已加载: {mp}")
|
||
except Exception:
|
||
print("加载失败，全新训练")
|
||
|
||
prompts = ['学而时习之', '床前明月光', '學而時習之', '春天', '人']
|
||
|
||
stages = [
|
||
(3e-3, 500, "快速下降"),
|
||
(1e-3, 500, "稳定收敛"),
|
||
(5e-4, 500, "精调"),
|
||
(2e-4, 500, "最终精调"),
|
||
]
|
||
|
||
t_total = time.time()
|
||
for lr, n_steps, label in stages:
|
||
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',
|
||
)
|
||
|
||
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,
|
||
)
|
||
|
||
print(f"\n完成，总耗时 {time.time()-t_total:.0f}s，已保存到 outputs/chat_model.pt")
|
||
|
||
|
||
if __name__ == '__main__':
|
||
main()
|