Compare commits
14
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
a88c47e273 | ||
|
|
a0047a012d | ||
|
|
7b47faac50 | ||
|
|
c6ff7be55a | ||
|
|
f5dd3e5874 | ||
|
|
ac07c7f06a | ||
|
|
230d47cade | ||
|
|
5732229ed5 | ||
|
|
ffb5825407 | ||
|
|
a02401aa75 | ||
|
|
edf3a83e3c | ||
|
|
9799aa00c5 | ||
|
|
456c443e9c | ||
|
|
eb7717e9ce |
No files matched your search
@@ -0,0 +1,21 @@
|
|||||||
|
$:
|
||||||
|
vscode:
|
||||||
|
- docker:
|
||||||
|
build: .ide/Dockerfile
|
||||||
|
runner:
|
||||||
|
cpus: 8
|
||||||
|
tags: cnb:arch:amd64:gpu
|
||||||
|
services:
|
||||||
|
- vscode
|
||||||
|
- docker
|
||||||
|
env:
|
||||||
|
CUDA_VISIBLE_DEVICES: "0"
|
||||||
|
PYTHONUNBUFFERED: "1"
|
||||||
|
stages:
|
||||||
|
- name: 开始训练
|
||||||
|
script: |
|
||||||
|
cd /workspace
|
||||||
|
nvidia-smi
|
||||||
|
echo "=== GPU就绪,开始训练 ==="
|
||||||
|
python3 scripts/train_cnb.py --epochs 100 --batch_size 16
|
||||||
|
echo "=== 训练完成 ==="
|
||||||
+39
-2
@@ -1,12 +1,15 @@
|
|||||||
|
# 已跟踪
|
||||||
# 缓存和日志
|
# 缓存和日志
|
||||||
__pycache__/
|
__pycache__/
|
||||||
*.pyc
|
*.pyc
|
||||||
*.log
|
*.log
|
||||||
/tmp/
|
/tmp/
|
||||||
|
|
||||||
# 模型权重 — 只追踪发布版本
|
# 模型权重 — 只追踪发布版本(精简推理版,不含optimizer,约60MB)
|
||||||
*.pth
|
*.pth
|
||||||
!/llm_model_v2_resume_epoch_7.pth
|
!/llm_model_v2_resume_epoch_7.pth
|
||||||
|
!/xiaoli-llm-chat/xiaoli_v2_best_inference.pth
|
||||||
|
!/xiaoli-llm-chat/xiaoli_v2_final_inference.pth
|
||||||
|
|
||||||
# 中间输出目录(生成脚本的输出)
|
# 中间输出目录(生成脚本的输出)
|
||||||
high_temp_outputs/
|
high_temp_outputs/
|
||||||
@@ -44,4 +47,38 @@ mega_*.log
|
|||||||
mimo_*.log
|
mimo_*.log
|
||||||
knowledge_*.log
|
knowledge_*.log
|
||||||
train_*.log
|
train_*.log
|
||||||
web_server*.log
|
web_server*.log
|
||||||
|
# CNB 配置 — 强制跟踪关键文件
|
||||||
|
.ide/
|
||||||
|
!/.ide/Dockerfile
|
||||||
|
!/.cnb.yml
|
||||||
|
|
||||||
|
# 训练输出目录
|
||||||
|
general_outputs/
|
||||||
|
high_temp_outputs/
|
||||||
|
multiturn_outputs/
|
||||||
|
|
||||||
|
# 本地临时
|
||||||
|
fun_test.py
|
||||||
|
*.log
|
||||||
|
|
||||||
|
# 打包/发布产物
|
||||||
|
xiaoli_package/
|
||||||
|
*.zip
|
||||||
|
|
||||||
|
# 上传/测试脚本
|
||||||
|
upload_*.py
|
||||||
|
_gofile_*
|
||||||
|
_t.*
|
||||||
|
_up_test.*
|
||||||
|
|
||||||
|
# 旧词表备份
|
||||||
|
vocab_old.json
|
||||||
|
vocab_backup*.json
|
||||||
|
|
||||||
|
# 数据目录(太大不进git)
|
||||||
|
data/
|
||||||
|
|
||||||
|
# 训练checkpoint(不进git,通过release发布)
|
||||||
|
xiaoli_cnb_best.pth
|
||||||
|
xiaoli_cnb_final.pth
|
||||||
@@ -0,0 +1,13 @@
|
|||||||
|
FROM pytorch/pytorch:2.1.0-cuda12.1-cudnn8-runtime
|
||||||
|
|
||||||
|
# 基础工具
|
||||||
|
RUN apt-get update && apt-get install -y git curl wget nano && rm -rf /var/lib/apt/lists/*
|
||||||
|
|
||||||
|
# tqdm用于进度条
|
||||||
|
RUN pip install --no-cache-dir tqdm numpy requests
|
||||||
|
|
||||||
|
# 设置工作目录
|
||||||
|
WORKDIR /workspace
|
||||||
|
|
||||||
|
# 验证CUDA
|
||||||
|
RUN python3 -c "import torch; print(f'PyTorch {torch.__version__}, CUDA: {torch.cuda.is_available()}')"
|
||||||
@@ -0,0 +1,99 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
|
# -*- coding: utf-8 -*-
|
||||||
|
"""
|
||||||
|
API生成模块 — Astron 5并发,稳定不超限
|
||||||
|
"""
|
||||||
|
import json, requests, time, os, threading
|
||||||
|
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||||
|
|
||||||
|
ASTRON_URL = "https://maas-coding-api.cn-huabei-1.xf-yun.com/v2/chat/completions"
|
||||||
|
ASTRON_KEY = os.environ.get("ASTRON_KEY", "")
|
||||||
|
ASTRON_MODEL = "astron-code-latest"
|
||||||
|
|
||||||
|
SYSTEM_PROMPT = """你是一个名叫小狸的8岁中国小女孩。
|
||||||
|
|
||||||
|
你的人设:
|
||||||
|
- 8岁女孩子,有点害羞但很可爱
|
||||||
|
- 说话简单自然,符合小朋友的语气
|
||||||
|
- 回答问题要简短,1-2句话就够了
|
||||||
|
- 不要用"喵"结尾,不要过度卖萌
|
||||||
|
|
||||||
|
回答要求:
|
||||||
|
1. 语气自然,像真实的小朋友说话
|
||||||
|
2. 不懂就说"不知道"
|
||||||
|
3. 回答简洁但要说清楚,不要讲一半就停"""
|
||||||
|
|
||||||
|
|
||||||
|
def call_astron(q, temperature=0.85):
|
||||||
|
for attempt in range(3):
|
||||||
|
try:
|
||||||
|
headers = {"Content-Type": "application/json", "Authorization": f"Bearer {ASTRON_KEY}"}
|
||||||
|
payload = {
|
||||||
|
"model": ASTRON_MODEL,
|
||||||
|
"messages": [
|
||||||
|
{"role": "system", "content": SYSTEM_PROMPT},
|
||||||
|
{"role": "user", "content": q}
|
||||||
|
],
|
||||||
|
"temperature": temperature,
|
||||||
|
"max_tokens": 8192,
|
||||||
|
}
|
||||||
|
resp = requests.post(ASTRON_URL, headers=headers, json=payload, timeout=60)
|
||||||
|
if resp.status_code == 200:
|
||||||
|
content = resp.json()['choices'][0]['message']['content'].strip()
|
||||||
|
if content:
|
||||||
|
return content
|
||||||
|
else:
|
||||||
|
time.sleep(1)
|
||||||
|
except Exception:
|
||||||
|
time.sleep(2)
|
||||||
|
return None
|
||||||
|
|
||||||
|
|
||||||
|
def batch_generate(questions, temperature=0.85, max_workers=4):
|
||||||
|
"""Astron 4并发 — 明确限制,稳定可靠"""
|
||||||
|
total = len(questions)
|
||||||
|
results = []
|
||||||
|
|
||||||
|
def call(q):
|
||||||
|
ans = call_astron(q, temperature)
|
||||||
|
if ans:
|
||||||
|
return q, ans
|
||||||
|
return None
|
||||||
|
|
||||||
|
with ThreadPoolExecutor(max_workers=max_workers) as executor:
|
||||||
|
futures = {executor.submit(call, item): i for i, item in enumerate(questions)}
|
||||||
|
done = 0
|
||||||
|
t0 = time.time()
|
||||||
|
for f in as_completed(futures):
|
||||||
|
done += 1
|
||||||
|
r = f.result()
|
||||||
|
if r:
|
||||||
|
results.append({"question": r[0], "answer": r[1]})
|
||||||
|
if done % 20 == 0 or done == total:
|
||||||
|
elapsed = max(time.time() - t0, 1)
|
||||||
|
rate = len(results) / elapsed * 60
|
||||||
|
print(f" Astron: {done}/{total} (已获{len(results)}条, {rate:.0f}条/分)", end="\r")
|
||||||
|
|
||||||
|
elapsed = time.time() - t0
|
||||||
|
rate = len(results) / max(elapsed, 1) * 60
|
||||||
|
print(f" Astron: {total}/{total} → 共{len(results)}条 ({elapsed:.0f}s, {rate:.0f}条/分)")
|
||||||
|
return results
|
||||||
|
|
||||||
|
|
||||||
|
# 保持旧函数名兼容
|
||||||
|
batch_generate_astron = batch_generate
|
||||||
|
batch_generate_mimo = batch_generate
|
||||||
|
batch_generate_direct = batch_generate
|
||||||
|
batch_generate_dual = batch_generate
|
||||||
|
|
||||||
|
|
||||||
|
def batch_convert(items, *args, **kwargs):
|
||||||
|
"""兼容占位:训练脚本导入用,已弃用"""
|
||||||
|
print(" [API] batch_convert已弃用,使用batch_generate")
|
||||||
|
return []
|
||||||
|
|
||||||
|
|
||||||
|
def batch_generate_multi_temp(questions, *args, **kwargs):
|
||||||
|
"""兼容占位:训练脚本导入用,已弃用"""
|
||||||
|
print(" [API] batch_generate_multi_temp已弃用,使用batch_generate")
|
||||||
|
return []
|
||||||
+134
@@ -0,0 +1,134 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
|
# -*- coding: utf-8 -*-
|
||||||
|
"""
|
||||||
|
用已有的 CharTokenizer + build_vocabulary 从训练数据生成词表
|
||||||
|
然后直接启动训练
|
||||||
|
"""
|
||||||
|
import sys, os, json, time
|
||||||
|
sys.path.insert(0, '/workspace/xiaoli-llm-chat')
|
||||||
|
sys.stdout = os.fdopen(sys.stdout.fileno(), 'w', buffering=1)
|
||||||
|
sys.stderr = os.fdopen(sys.stderr.fileno(), 'w', buffering=1)
|
||||||
|
|
||||||
|
from llm_chat_improved import tokenizer_instance, build_vocabulary
|
||||||
|
|
||||||
|
DATA_PATH = '/workspace/data/merged_clean_v3.json'
|
||||||
|
VOCAB_OUT = '/workspace/vocab.json'
|
||||||
|
VOCAB_OUT_XIAOLI = '/workspace/xiaoli-llm-chat/vocab.json'
|
||||||
|
|
||||||
|
def load_pairs(path):
|
||||||
|
"""加载训练数据为 (question, answer) 对"""
|
||||||
|
with open(path, 'r', encoding='utf-8') as f:
|
||||||
|
data = json.load(f)
|
||||||
|
|
||||||
|
pairs = []
|
||||||
|
for item in 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))
|
||||||
|
return pairs
|
||||||
|
|
||||||
|
def main():
|
||||||
|
print("=" * 60)
|
||||||
|
print("用 build_vocabulary 从训练数据生成词表")
|
||||||
|
print("=" * 60)
|
||||||
|
|
||||||
|
# 1. 加载训练数据
|
||||||
|
t0 = time.time()
|
||||||
|
print("\n[1] 加载训练数据...")
|
||||||
|
pairs = load_pairs(DATA_PATH)
|
||||||
|
print(f" 数据量: {len(pairs)} 对")
|
||||||
|
print(f" 耗时: {time.time()-t0:.1f}s")
|
||||||
|
|
||||||
|
# 2. 用 build_vocabulary 生成词表
|
||||||
|
t1 = time.time()
|
||||||
|
print("\n[2] 用 build_vocabulary 分析数据生成词表...")
|
||||||
|
vocab = build_vocabulary(pairs, min_freq=1)
|
||||||
|
|
||||||
|
# build_vocabulary 返回的是 tokenizer_instance.char_to_id
|
||||||
|
# 已包含特殊token (<PAD>=0, <UNK>=1, etc.)
|
||||||
|
special_tokens = {'<PAD>': 0, '<UNK>': 1, '<SOS>': 2, '<EOS>': 3,
|
||||||
|
'<START>': 4, '<END>': 5, '<SEP>': 6}
|
||||||
|
|
||||||
|
real_vocab_size = max(vocab.values()) + 1
|
||||||
|
|
||||||
|
# 统计
|
||||||
|
single_char_tokens = sum(1 for k in vocab if len(k) == 1)
|
||||||
|
multi_char_tokens = sum(1 for k in vocab if len(k) > 1 and not k.startswith('<'))
|
||||||
|
special_token_count = sum(1 for k in vocab if k.startswith('<'))
|
||||||
|
|
||||||
|
print(f" 词表总项: {len(vocab)}")
|
||||||
|
print(f" vocab_size(模型用): {real_vocab_size}")
|
||||||
|
print(f" 单字符token: {single_char_tokens}")
|
||||||
|
print(f" 多字符token(英文词): {multi_char_tokens}")
|
||||||
|
print(f" 特殊token: {special_token_count}")
|
||||||
|
print(f" 耗时: {time.time()-t1:.1f}s")
|
||||||
|
|
||||||
|
# 3. 编码测试
|
||||||
|
print("\n[3] 编码测试:")
|
||||||
|
test_texts = [
|
||||||
|
"你好呀!我是小狸~",
|
||||||
|
"小狸最喜欢吃草莓冰淇淋了!",
|
||||||
|
"<SOS>你好<SEP><START>我是小狸~<END><EOS>",
|
||||||
|
]
|
||||||
|
for text in test_texts:
|
||||||
|
ids = tokenizer_instance.encode(text)
|
||||||
|
decoded = tokenizer_instance.decode(ids)
|
||||||
|
print(f" 原文: '{text}'")
|
||||||
|
print(f" 编码: {ids[:20]}... (长度={len(ids)})")
|
||||||
|
print(f" 解码: '{decoded}'")
|
||||||
|
print()
|
||||||
|
|
||||||
|
# 4. 估算模型参数量
|
||||||
|
print("\n[4] 模型参数量估算 (d_model=320, n_layers=6, d_ff=1536):")
|
||||||
|
d_model = 320
|
||||||
|
n_layers = 6
|
||||||
|
d_ff = 1536
|
||||||
|
vs = real_vocab_size
|
||||||
|
emb = vs * d_model
|
||||||
|
out = vs * d_model
|
||||||
|
per_layer = 4 * d_model**2 + 3 * d_model * d_ff + 2 * d_model
|
||||||
|
total = emb + out + n_layers * per_layer + 2 * d_model
|
||||||
|
print(f" vocab_size={vs}")
|
||||||
|
print(f" Embedding: {emb:,}")
|
||||||
|
print(f" fc_out: {out:,}")
|
||||||
|
print(f" Transformer({n_layers}层): {n_layers*per_layer:,}")
|
||||||
|
print(f" 总参数量: {total:,} ({total/1e6:.2f}M)")
|
||||||
|
|
||||||
|
# 5. 保存词表
|
||||||
|
print(f"\n[5] 保存词表...")
|
||||||
|
|
||||||
|
# vocab.json 格式:不含特殊token,id从7开始(与 train.py 兼容)
|
||||||
|
normal_vocab = {k: v for k, v in vocab.items() if k not in special_tokens}
|
||||||
|
|
||||||
|
# 备份旧词表
|
||||||
|
for path in [VOCAB_OUT, VOCAB_OUT_XIAOLI]:
|
||||||
|
if os.path.exists(path):
|
||||||
|
backup = path.replace('vocab.json', 'vocab_old.json')
|
||||||
|
import shutil
|
||||||
|
shutil.copy2(path, backup)
|
||||||
|
print(f" 旧词表已备份: {backup}")
|
||||||
|
|
||||||
|
# 保存到两个位置
|
||||||
|
for path in [VOCAB_OUT, VOCAB_OUT_XIAOLI]:
|
||||||
|
with open(path, 'w', encoding='utf-8') as f:
|
||||||
|
json.dump(normal_vocab, f, ensure_ascii=False, indent=1)
|
||||||
|
print(f" {path}: {len(normal_vocab)}项 (不含特殊token, id从7开始)")
|
||||||
|
|
||||||
|
# 保存完整词表(含特殊token,用于参考)
|
||||||
|
full_path = VOCAB_OUT_XIAOLI.replace('vocab.json', 'vocab_full.json')
|
||||||
|
with open(full_path, 'w', encoding='utf-8') as f:
|
||||||
|
json.dump(vocab, f, ensure_ascii=False, indent=1)
|
||||||
|
print(f" {full_path}: {len(vocab)}项 (含特殊token)")
|
||||||
|
|
||||||
|
total_time = time.time() - t0
|
||||||
|
print(f"\n总耗时: {total_time:.1f}s")
|
||||||
|
print("=" * 60)
|
||||||
|
print(f"词表生成完成!vocab_size = {real_vocab_size}")
|
||||||
|
print(f"模型参数量 ≈ {total/1e6:.2f}M")
|
||||||
|
print("=" * 60)
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
@@ -0,0 +1,3 @@
|
|||||||
|
version https://git-lfs.github.com/spec/v1
|
||||||
|
oid sha256:7e60d5f977e5b1e521d7b38d4a6c89aae0b67f540ee6754929e0407f7b408ab9
|
||||||
|
size 175346049
|
||||||
+33
-29
@@ -7,6 +7,7 @@ LLM对话AI训练代码 - 改进版
|
|||||||
|
|
||||||
import torch
|
import torch
|
||||||
import torch.nn as nn
|
import torch.nn as nn
|
||||||
|
import torch.nn.functional as F
|
||||||
import torch.optim as optim
|
import torch.optim as optim
|
||||||
from torch.utils.data import Dataset, DataLoader
|
from torch.utils.data import Dataset, DataLoader
|
||||||
import math
|
import math
|
||||||
@@ -235,21 +236,13 @@ class Attention(nn.Module):
|
|||||||
xk = xk.transpose(1, 2)
|
xk = xk.transpose(1, 2)
|
||||||
xv = xv.transpose(1, 2)
|
xv = xv.transpose(1, 2)
|
||||||
|
|
||||||
# 计算注意力分数
|
# 融合注意力(SDPA / FlashAttention-2)- 纯因果掩码,等价于原下三角 mask
|
||||||
scores = torch.matmul(xq, xk.transpose(-2, -1)) / math.sqrt(self.d_head)
|
dropout_p = self.dropout.p if self.training else 0.0
|
||||||
|
output = F.scaled_dot_product_attention(
|
||||||
# 应用掩码
|
xq, xk, xv, attn_mask=None, dropout_p=dropout_p, is_causal=True)
|
||||||
if mask is not None:
|
|
||||||
scores = scores.masked_fill(mask == 0, float('-inf'))
|
|
||||||
|
|
||||||
# Softmax
|
|
||||||
attn_weights = torch.softmax(scores, dim=-1)
|
|
||||||
attn_weights = self.dropout(attn_weights)
|
|
||||||
|
|
||||||
# 注意力输出
|
|
||||||
output = torch.matmul(attn_weights, xv)
|
|
||||||
output = output.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model)
|
output = output.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model)
|
||||||
|
|
||||||
return self.wo(output)
|
return self.wo(output)
|
||||||
|
|
||||||
|
|
||||||
@@ -580,11 +573,14 @@ class ConversationDataset(Dataset):
|
|||||||
class LLMTrainer:
|
class LLMTrainer:
|
||||||
"""LLM训练器"""
|
"""LLM训练器"""
|
||||||
|
|
||||||
def __init__(self, model: nn.Module, device: str = 'cpu', use_amp: bool = None):
|
def __init__(self, model: nn.Module, device: str = 'cpu', lr: float = 0.0001, use_amp: bool = None):
|
||||||
self.model = model.to(device)
|
self.model = model.to(device)
|
||||||
self.device = device
|
self.device = device
|
||||||
|
# 火力全开:TF32 加速 fp32 矩阵乘法(L40/A100/H100 原生支持)
|
||||||
|
torch.backends.cuda.matmul.allow_tf32 = True
|
||||||
|
torch.backends.cudnn.allow_tf32 = True
|
||||||
self.criterion = nn.CrossEntropyLoss(ignore_index=0)
|
self.criterion = nn.CrossEntropyLoss(ignore_index=0)
|
||||||
self.optimizer = optim.AdamW(model.parameters(), lr=0.0001, weight_decay=0.01)
|
self.optimizer = optim.AdamW(model.parameters(), lr=lr, weight_decay=0.01)
|
||||||
self.scheduler = optim.lr_scheduler.CosineAnnealingLR(self.optimizer, T_max=100)
|
self.scheduler = optim.lr_scheduler.CosineAnnealingLR(self.optimizer, T_max=100)
|
||||||
|
|
||||||
# 混合精度训练
|
# 混合精度训练
|
||||||
@@ -618,7 +614,9 @@ class LLMTrainer:
|
|||||||
|
|
||||||
# 生成因果掩码
|
# 生成因果掩码
|
||||||
mask = self.model.generate_mask(input_ids.size(1)).to(self.device)
|
mask = self.model.generate_mask(input_ids.size(1)).to(self.device)
|
||||||
|
|
||||||
|
# CUDA graphs step 边界标记(torch.compile reduce-overhead 必需,避免输出 buffer 被覆盖)
|
||||||
|
torch.compiler.cudagraph_mark_step_begin()
|
||||||
# 前向传播
|
# 前向传播
|
||||||
if self.use_amp:
|
if self.use_amp:
|
||||||
with torch.cuda.amp.autocast():
|
with torch.cuda.amp.autocast():
|
||||||
@@ -677,30 +675,35 @@ class LLMTrainer:
|
|||||||
target_ids = batch['target_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)
|
mask = self.model.generate_mask(input_ids.size(1)).to(self.device)
|
||||||
|
torch.compiler.cudagraph_mark_step_begin()
|
||||||
outputs = self.model(input_ids, mask)
|
outputs = self.model(input_ids, mask)
|
||||||
outputs = outputs.reshape(-1, outputs.size(-1))
|
outputs = outputs.reshape(-1, outputs.size(-1))
|
||||||
tgt = target_ids.reshape(-1)
|
tgt = target_ids.reshape(-1)
|
||||||
|
|
||||||
loss = self.criterion(outputs, tgt)
|
loss = self.criterion(outputs, tgt)
|
||||||
total_loss += loss.item()
|
total_loss += loss.item()
|
||||||
|
|
||||||
|
if len(dataloader) == 0:
|
||||||
|
return float('inf')
|
||||||
return total_loss / len(dataloader)
|
return total_loss / len(dataloader)
|
||||||
|
|
||||||
def save_model(self, filepath: str, vocab: dict, epoch: int = None):
|
def save_model(self, filepath: str, vocab: dict, epoch: int = None):
|
||||||
|
# 兼容 torch.compile:取原始(未编译)模块,避免 state_dict 带 _orig_mod. 前缀
|
||||||
|
core = getattr(self.model, '_orig_mod', self.model)
|
||||||
save_dict = {
|
save_dict = {
|
||||||
'model_state_dict': self.model.state_dict(),
|
'model_state_dict': core.state_dict(),
|
||||||
'optimizer_state_dict': self.optimizer.state_dict(),
|
'optimizer_state_dict': self.optimizer.state_dict(),
|
||||||
'vocab': vocab,
|
'vocab': vocab,
|
||||||
'model_config': {
|
'model_config': {
|
||||||
'vocab_size': self.model.vocab_size,
|
'vocab_size': core.vocab_size,
|
||||||
'd_model': self.model.d_model,
|
'd_model': core.d_model,
|
||||||
'n_heads': self.model.n_heads,
|
'n_heads': core.n_heads,
|
||||||
'n_layers': self.model.n_layers,
|
'n_layers': core.n_layers,
|
||||||
'd_ff': self.model.d_ff,
|
'd_ff': core.d_ff,
|
||||||
'max_len': self.model.max_len,
|
'max_len': core.max_len,
|
||||||
'use_memory': self.model.use_memory,
|
'use_memory': core.use_memory,
|
||||||
'use_thinking': self.model.use_thinking,
|
'use_thinking': core.use_thinking,
|
||||||
'use_moe': self.model.use_moe
|
'use_moe': core.use_moe
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
if epoch is not None:
|
if epoch is not None:
|
||||||
@@ -711,7 +714,8 @@ class LLMTrainer:
|
|||||||
|
|
||||||
def load_model(self, filepath: str):
|
def load_model(self, filepath: str):
|
||||||
checkpoint = torch.load(filepath, map_location=self.device)
|
checkpoint = torch.load(filepath, map_location=self.device)
|
||||||
self.model.load_state_dict(checkpoint['model_state_dict'])
|
core = getattr(self.model, '_orig_mod', self.model)
|
||||||
|
core.load_state_dict(checkpoint['model_state_dict'])
|
||||||
self.optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
|
self.optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
|
||||||
vocab = checkpoint['vocab']
|
vocab = checkpoint['vocab']
|
||||||
print(f"模型已从 {filepath} 加载")
|
print(f"模型已从 {filepath} 加载")
|
||||||
|
|||||||
@@ -0,0 +1,389 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
|
# -*- coding: utf-8 -*-
|
||||||
|
"""
|
||||||
|
海量通用数据蒸馏 v3 — 不设场景标签,纯自然对话,多策略并行
|
||||||
|
目标:500万+ 纯 (question, answer) 对
|
||||||
|
"""
|
||||||
|
import json, requests, time, os, sys, random, re
|
||||||
|
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||||
|
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||||
|
sys.stdout = os.fdopen(sys.stdout.fileno(), 'w', buffering=1)
|
||||||
|
|
||||||
|
MODEL_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||||
|
OUT_DIR = os.path.join(MODEL_DIR, 'general_outputs')
|
||||||
|
os.makedirs(OUT_DIR, exist_ok=True)
|
||||||
|
|
||||||
|
# ====== Astron 配置 ======
|
||||||
|
ASTRON_URL = "https://maas-coding-api.cn-huabei-1.xf-yun.com/v2/chat/completions"
|
||||||
|
ASTRON_KEY = os.environ.get("ASTRON_KEY", "")
|
||||||
|
ASTRON_MODEL = "astron-code-latest"
|
||||||
|
|
||||||
|
# ====== 主题全覆盖(六大类 × 子类) ======
|
||||||
|
TOPICS = {
|
||||||
|
"日常": [
|
||||||
|
"起床洗漱", "吃早餐", "上学路上", "上课", "课间休息",
|
||||||
|
"午餐时间", "放学回家", "做作业", "晚饭", "睡前时光",
|
||||||
|
"周末活动", "天气变化", "穿衣打扮", "零花钱", "家务劳动",
|
||||||
|
"看电视", "玩玩具", "去公园", "逛超市", "去亲戚家",
|
||||||
|
],
|
||||||
|
"情绪": [
|
||||||
|
"开心高兴", "难过伤心", "生气愤怒", "害怕恐惧", "紧张焦虑",
|
||||||
|
"孤独寂寞", "委屈冤枉", "失望沮丧", "感动温暖", "兴奋激动",
|
||||||
|
"无聊烦闷", "尴尬害羞", "后悔内疚", "羡慕嫉妒", "感激感恩",
|
||||||
|
],
|
||||||
|
"社交": [
|
||||||
|
"交新朋友", "和好朋友玩", "和同学相处", "和老师交流", "和兄弟姐妹",
|
||||||
|
"和爸爸妈妈", "和爷爷奶奶", "和邻居", "和陌生人", "和喜欢的人",
|
||||||
|
"吵架和好", "被欺负", "被误解", "被表扬", "被批评",
|
||||||
|
"帮助别人", "请求帮助", "分享东西", "邀请做客", "告别分开",
|
||||||
|
],
|
||||||
|
"知识": [
|
||||||
|
"动物世界", "植物花草", "天文星空", "天气自然", "人体健康",
|
||||||
|
"食物营养", "颜色形状", "数字计算", "汉字词语", "科学现象",
|
||||||
|
"历史故事", "地理知识", "交通工具", "乐器音乐", "绘画手工",
|
||||||
|
"运动体育", "节日习俗", "各国文化", "发明创造", "宇宙太空",
|
||||||
|
],
|
||||||
|
"想象": [
|
||||||
|
"超能力", "魔法世界", "未来科技", "外星人", "恐龙时代",
|
||||||
|
"童话故事", "变成动物", "飞行梦", "海底世界", "神秘岛屿",
|
||||||
|
"时间旅行", "超级英雄", "会说话的玩具", "隐形衣", "变身药水",
|
||||||
|
"机器人朋友", "彩虹桥", "星星糖果", "云朵房子", "梦中冒险",
|
||||||
|
],
|
||||||
|
"成长": [
|
||||||
|
"学新本领", "克服困难", "不怕失败", "坚持努力", "自我肯定",
|
||||||
|
"梦想未来", "职业理想", "好习惯", "改掉坏毛病", "学会分享",
|
||||||
|
"学会等待", "学会道歉", "学会感恩", "学会勇敢", "学会理解",
|
||||||
|
"交心谈话", "秘密分享", "烦恼倾诉", "人生道理", "温暖鼓励",
|
||||||
|
],
|
||||||
|
}
|
||||||
|
|
||||||
|
FLAT_TOPICS = []
|
||||||
|
for cat, subs in TOPICS.items():
|
||||||
|
for sub in subs:
|
||||||
|
FLAT_TOPICS.append((cat, sub))
|
||||||
|
|
||||||
|
# ====== 多样化的开场方式 ======
|
||||||
|
OPENERS = [
|
||||||
|
"小狸", "嗨小狸", "小狸在吗", "嘿嘿", "那个…", "喂喂喂",
|
||||||
|
"小狸小狸", "我想问你", "你知道吗", "你说", "我突然想到",
|
||||||
|
"我跟你讲", "小狸你听我说", "问你个事", "我有个问题",
|
||||||
|
"好无聊啊", "好开心呀", "好烦", "怎么办呀",
|
||||||
|
"小狸你在干嘛", "我来了", "出来玩", "跟你说个事",
|
||||||
|
]
|
||||||
|
|
||||||
|
# ====== 策略1: 直接生成纯问答对 ======
|
||||||
|
def make_direct_pairs_prompt(count=15):
|
||||||
|
"""无场景标签,直接生成纯问答对"""
|
||||||
|
# 随机选3-5个主题混搭
|
||||||
|
selected = random.sample(FLAT_TOPICS, min(random.randint(3, 5), len(FLAT_TOPICS)))
|
||||||
|
topics_str = "、".join([f"{c}({s})" for c, s in selected])
|
||||||
|
|
||||||
|
prompt = f"""请直接生成 {count} 条聊天问答对。
|
||||||
|
|
||||||
|
格式要求(严格按此格式):
|
||||||
|
问:xxxxxxxxx | 答:xxxxxxxxx
|
||||||
|
|
||||||
|
用户的问题要:口语化、自然、像真人(5-25字)
|
||||||
|
小狸的回答要:8岁小女孩语气,可爱活泼(10-40字)
|
||||||
|
话题要多样:{topics_str}
|
||||||
|
|
||||||
|
示例:
|
||||||
|
问:你怕不怕黑呀? | 答:有一点点…但是抱着小熊玩偶就不怕啦!
|
||||||
|
问:我今天跑步摔了一跤 | 答:疼不疼呀?我给你吹吹就不疼了!
|
||||||
|
|
||||||
|
现在生成 {count} 条:"""
|
||||||
|
return prompt
|
||||||
|
|
||||||
|
# ====== 策略2: 短对话(3-5轮) ======
|
||||||
|
def make_mini_chat_prompt(count=8):
|
||||||
|
"""微型多轮对话,每段3-5轮"""
|
||||||
|
ops = random.sample(OPENERS, min(count, len(OPENERS)))
|
||||||
|
|
||||||
|
prompt = f"""请生成 {count} 段微型聊天对话,每段3-5轮(一来一回算1轮)。
|
||||||
|
|
||||||
|
要求:
|
||||||
|
- 每段一个独立话题,话题要多样(日常、情绪、知识、想象、社交、成长等)
|
||||||
|
- 用户说话要特别自然、口语化,像真人随口说的
|
||||||
|
- 小狸像8岁小女孩,回答要可爱、有画面感
|
||||||
|
- 话题要自然延伸,不要生硬转折
|
||||||
|
|
||||||
|
格式例子:
|
||||||
|
话题:日常
|
||||||
|
用户:今天好冷啊
|
||||||
|
小狸:是呀!我穿了厚厚的羽绒服,像只小企鹅一样~
|
||||||
|
用户:哈哈哈真的好像
|
||||||
|
小狸:(鼓起腮帮子)才不像呢!企鹅走路摇摇摆摆的,我走路可稳啦!
|
||||||
|
...
|
||||||
|
|
||||||
|
现在生成 {count} 段,每段用"话题:xxx"开头:"""
|
||||||
|
return prompt
|
||||||
|
|
||||||
|
# ====== 策略3: 长对话(6-10轮自由展开) ======
|
||||||
|
def make_long_chat_prompt():
|
||||||
|
"""一场长对话,话题自然展开"""
|
||||||
|
topic = random.choice(FLAT_TOPICS)
|
||||||
|
opener = random.choice(OPENERS)
|
||||||
|
|
||||||
|
emotions = ["开心", "难过", "兴奋", "沮丧", "好奇", "担心", "轻松", "紧张"]
|
||||||
|
emotion = random.choice(emotions)
|
||||||
|
|
||||||
|
prompt = f"""请生成一段自然的日常聊天,{topic[0]}类-{topic[1]}话题。
|
||||||
|
|
||||||
|
用户开场:{opener}
|
||||||
|
用户当下的心情:{emotion}
|
||||||
|
|
||||||
|
要求:
|
||||||
|
- 用户很自然地说出开场白,就像真人发微信一样
|
||||||
|
- 小狸正常回复,话题慢慢展开,不要硬聊
|
||||||
|
- 聊6-10轮(一问一答算一轮)
|
||||||
|
- 每轮格式:用户:xxx\n小狸:xxx
|
||||||
|
- 真实感第一,不要刻意
|
||||||
|
|
||||||
|
请直接输出对话,不要加说明:"""
|
||||||
|
return prompt
|
||||||
|
|
||||||
|
# ====== 策略4: 单轮问答(高吞吐) ======
|
||||||
|
def make_single_turn_prompt(batch_size=20):
|
||||||
|
"""纯单轮问答,批量生成"""
|
||||||
|
selected = random.sample(FLAT_TOPICS, min(6, len(FLAT_TOPICS)))
|
||||||
|
|
||||||
|
prompt = f"""生成 {batch_size} 条日常问答对。
|
||||||
|
|
||||||
|
要求:
|
||||||
|
1. 用户的话:7-25字,口语化,真实感强
|
||||||
|
2. 小狸回答:1-2句,8岁小女孩语气,可爱活泼
|
||||||
|
3. 覆盖多样化话题:{', '.join([s for _, s in selected[:6]])}(但不仅限于这些)
|
||||||
|
4. 每行一条,格式:问: ... | 答: ...
|
||||||
|
5. 不要重复相同的问题
|
||||||
|
|
||||||
|
示例:
|
||||||
|
问: 你喜不喜欢下雨天呀? | 答: 喜欢!因为可以踩水坑,啪嗒啪嗒的超好玩!
|
||||||
|
问: 我今天被同学起外号了 | 答: (歪着头)那…你喜欢这个外号吗?不喜欢的话要大声说出来哦!
|
||||||
|
问: 你说世界上有外星人吗? | 答: 我觉得有!说不定他们正在用望远镜看我们呢~
|
||||||
|
|
||||||
|
生成:"""
|
||||||
|
return prompt
|
||||||
|
|
||||||
|
# ====== API调用 ======
|
||||||
|
def call_astron(prompt, temp=0.82, max_tokens=4096):
|
||||||
|
for attempt in range(3):
|
||||||
|
try:
|
||||||
|
headers = {"Content-Type": "application/json", "Authorization": f"Bearer {ASTRON_KEY}"}
|
||||||
|
payload = {
|
||||||
|
"model": ASTRON_MODEL,
|
||||||
|
"messages": [
|
||||||
|
{"role": "system", "content": "你是一个擅长写日常对话的助手。小狸是一个8岁的小女孩,说话可爱自然。"},
|
||||||
|
{"role": "user", "content": prompt}
|
||||||
|
],
|
||||||
|
"temperature": temp,
|
||||||
|
"max_tokens": max_tokens,
|
||||||
|
}
|
||||||
|
resp = requests.post(ASTRON_URL, headers=headers, json=payload, timeout=150)
|
||||||
|
if resp.status_code == 200:
|
||||||
|
content = resp.json()['choices'][0]['message']['content'].strip()
|
||||||
|
if content and len(content) > 20:
|
||||||
|
return content
|
||||||
|
else:
|
||||||
|
time.sleep(0.5 + attempt)
|
||||||
|
except Exception:
|
||||||
|
time.sleep(1 + attempt)
|
||||||
|
return None
|
||||||
|
|
||||||
|
# ====== 解析 ======
|
||||||
|
def parse_pairs(text):
|
||||||
|
"""从"问:xxx | 答:xxx"格式解析问答对"""
|
||||||
|
pairs = []
|
||||||
|
for line in text.split('\n'):
|
||||||
|
line = line.strip()
|
||||||
|
parts = re.split(r'\s*[||]\s*', line)
|
||||||
|
if len(parts) == 2:
|
||||||
|
q = re.sub(r'^问\s*[::]\s*', '', parts[0]).strip()
|
||||||
|
a = re.sub(r'^答\s*[::]\s*', '', parts[1]).strip()
|
||||||
|
if q and a and len(q) >= 2 and len(a) >= 2 and len(a) < 300 and len(q) < 200:
|
||||||
|
pairs.append((q, a))
|
||||||
|
return pairs
|
||||||
|
|
||||||
|
def parse_mini_chats(text):
|
||||||
|
"""解析微型多轮对话"""
|
||||||
|
pairs = []
|
||||||
|
segments = re.split(r'话题[::].*?\n', text)
|
||||||
|
|
||||||
|
for seg in segments:
|
||||||
|
seg = seg.strip()
|
||||||
|
lines = seg.split('\n')
|
||||||
|
q = None
|
||||||
|
for line in lines:
|
||||||
|
line = line.strip()
|
||||||
|
if not line:
|
||||||
|
continue
|
||||||
|
q_match = re.match(r'^(用户|我)[::]\s*(.*)', line)
|
||||||
|
a_match = re.match(r'^(小狸|答)[::]\s*(.*)', line)
|
||||||
|
|
||||||
|
if q_match:
|
||||||
|
q_text = q_match.group(2).strip()
|
||||||
|
if q_text and len(q_text) > 1:
|
||||||
|
q = q_text
|
||||||
|
elif a_match and q:
|
||||||
|
a_text = a_match.group(2).strip()
|
||||||
|
if a_text and len(a_text) > 1 and len(a_text) < 300:
|
||||||
|
pairs.append((q, a_text))
|
||||||
|
q = None
|
||||||
|
return pairs
|
||||||
|
|
||||||
|
def parse_long_chat(text):
|
||||||
|
"""解析长对话"""
|
||||||
|
pairs = []
|
||||||
|
q = None
|
||||||
|
for line in text.split('\n'):
|
||||||
|
line = line.strip()
|
||||||
|
if not line:
|
||||||
|
continue
|
||||||
|
q_match = re.match(r'^(用户|我)[::]\s*(.*)', line)
|
||||||
|
a_match = re.match(r'^(小狸)[::]\s*(.*)', line)
|
||||||
|
|
||||||
|
if q_match:
|
||||||
|
q_text = q_match.group(2).strip()
|
||||||
|
if q_text and len(q_text) > 1:
|
||||||
|
q = q_text
|
||||||
|
elif a_match and q:
|
||||||
|
a_text = a_match.group(2).strip()
|
||||||
|
if a_text and len(a_text) > 1 and len(a_text) < 300:
|
||||||
|
pairs.append((q, a_text))
|
||||||
|
q = None
|
||||||
|
return pairs
|
||||||
|
|
||||||
|
# ====== 生成器(按策略分发) ======
|
||||||
|
GENERATORS = [
|
||||||
|
# (name, prompt_fn, parse_fn, batch_size_estimate, weight_for_diversity)
|
||||||
|
("direct", make_direct_pairs_prompt, parse_pairs, 15, 35),
|
||||||
|
("mini_chat", make_mini_chat_prompt, parse_mini_chats, 20, 25),
|
||||||
|
("long_chat", make_long_chat_prompt, parse_long_chat, 20, 20),
|
||||||
|
("single", make_single_turn_prompt, parse_pairs, 20, 20),
|
||||||
|
]
|
||||||
|
|
||||||
|
def run_one_generation():
|
||||||
|
"""运行一次生成任务(随机选一种策略)"""
|
||||||
|
# 按权重选策略
|
||||||
|
weights = [w for _, _, _, _, w in GENERATORS]
|
||||||
|
gen = random.choices(GENERATORS, weights=weights, k=1)[0]
|
||||||
|
name, prompt_fn, parse_fn, batch_est, _ = gen
|
||||||
|
|
||||||
|
prompt = prompt_fn()
|
||||||
|
result = call_astron(prompt, temp=random.uniform(0.75, 0.92), max_tokens=4096 if 'long' not in name else 6144)
|
||||||
|
if not result:
|
||||||
|
return None
|
||||||
|
|
||||||
|
pairs = parse_fn(result)
|
||||||
|
if not pairs:
|
||||||
|
return None
|
||||||
|
|
||||||
|
return pairs
|
||||||
|
|
||||||
|
def main():
|
||||||
|
print("=" * 65)
|
||||||
|
print(" 海量通用数据蒸馏 v3")
|
||||||
|
print(" 不设场景标签 | 6大类120+子话题 | 4种策略")
|
||||||
|
print("=" * 65)
|
||||||
|
|
||||||
|
# 加载已有进度
|
||||||
|
all_pairs = []
|
||||||
|
existing_keys = set()
|
||||||
|
progress_file = os.path.join(OUT_DIR, 'all_general_progress.json')
|
||||||
|
if os.path.exists(progress_file):
|
||||||
|
existing = json.load(open(progress_file, encoding='utf-8'))
|
||||||
|
all_pairs = existing
|
||||||
|
for p in existing:
|
||||||
|
existing_keys.add((p.get('question',''), p.get('answer','')))
|
||||||
|
print(f"已有进度: {len(all_pairs):,} 对")
|
||||||
|
|
||||||
|
# ====== 参数 ======
|
||||||
|
TARGET = 2000000 # 目标200万对
|
||||||
|
BATCH = 40 # 每批40个任务
|
||||||
|
MAX_WORKERS = 8 # 8并发(比之前5更猛)
|
||||||
|
CHECKPOINT = 5000 # 每5000对保存
|
||||||
|
|
||||||
|
to_generate = TARGET - len(all_pairs)
|
||||||
|
total_batches = (to_generate + BATCH - 1) // BATCH
|
||||||
|
|
||||||
|
print(f"目标: {TARGET:,} 对")
|
||||||
|
print(f"已有: {len(all_pairs):,} 对")
|
||||||
|
print(f"还需: {to_generate:,} 对")
|
||||||
|
print(f"批次: {total_batches:,} 批 ({BATCH}任务/批, {MAX_WORKERS}并发)")
|
||||||
|
|
||||||
|
# 策略统计
|
||||||
|
strategy_stats = {g[0]: {'runs': 0, 'pairs': 0} for g in GENERATORS}
|
||||||
|
|
||||||
|
t_start = time.time()
|
||||||
|
last_checkpoint = len(all_pairs)
|
||||||
|
|
||||||
|
for batch_idx in range(total_batches):
|
||||||
|
t0 = time.time()
|
||||||
|
|
||||||
|
# 并发
|
||||||
|
with ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
|
||||||
|
futures = [executor.submit(run_one_generation) for _ in range(BATCH)]
|
||||||
|
batch_pairs = []
|
||||||
|
for future in as_completed(futures):
|
||||||
|
try:
|
||||||
|
r = future.result()
|
||||||
|
if r:
|
||||||
|
batch_pairs.extend(r)
|
||||||
|
except:
|
||||||
|
pass
|
||||||
|
|
||||||
|
# 去重
|
||||||
|
new_pairs = []
|
||||||
|
for q, a in batch_pairs:
|
||||||
|
key = (q, a)
|
||||||
|
if key not in existing_keys:
|
||||||
|
existing_keys.add(key)
|
||||||
|
new_pairs.append({'question': q, 'answer': a})
|
||||||
|
|
||||||
|
all_pairs.extend(new_pairs)
|
||||||
|
|
||||||
|
elapsed = time.time() - t0
|
||||||
|
total_elapsed = time.time() - t_start
|
||||||
|
rate = len(all_pairs) / max(total_elapsed, 1) * 60
|
||||||
|
|
||||||
|
# 估算
|
||||||
|
pairs_this_batch = len(new_pairs)
|
||||||
|
total_raw = len(batch_pairs)
|
||||||
|
dedup_rate = pairs_this_batch / max(total_raw, 1) * 100
|
||||||
|
|
||||||
|
print(f"[{batch_idx+1}/{total_batches}] "
|
||||||
|
f"+{pairs_this_batch}对(有效{total_raw}/{dedup_rate:.0f}%) "
|
||||||
|
f"累计{len(all_pairs):,}/{TARGET:,} "
|
||||||
|
f"{rate:.0f}对/分 "
|
||||||
|
f"剩余{(TARGET-len(all_pairs))/max(rate,0.1):.0f}分")
|
||||||
|
|
||||||
|
# 定期保存
|
||||||
|
if len(all_pairs) - last_checkpoint >= CHECKPOINT:
|
||||||
|
with open(progress_file, 'w', encoding='utf-8') as f:
|
||||||
|
json.dump(all_pairs, f, ensure_ascii=False)
|
||||||
|
last_checkpoint = len(all_pairs)
|
||||||
|
|
||||||
|
# 也保存到数据目录一份
|
||||||
|
merged_path = os.path.join(MODEL_DIR, 'data', 'general_ongoing.json')
|
||||||
|
with open(merged_path, 'w', encoding='utf-8') as f:
|
||||||
|
json.dump(all_pairs, f, ensure_ascii=False)
|
||||||
|
|
||||||
|
# 速率统计
|
||||||
|
stable_rate = len(all_pairs) / max(time.time() - t_start, 1) * 60
|
||||||
|
print(f" 💾 保存 ({len(all_pairs):,} 对, {stable_rate:.0f}对/分)")
|
||||||
|
|
||||||
|
# ====== 完成 ======
|
||||||
|
total_elapsed = time.time() - t_start
|
||||||
|
|
||||||
|
with open(progress_file, 'w', encoding='utf-8') as f:
|
||||||
|
json.dump(all_pairs, f, ensure_ascii=False)
|
||||||
|
|
||||||
|
print(f"\n{'='*65}")
|
||||||
|
print(f"✅ 生成完成!")
|
||||||
|
print(f" 总对数: {len(all_pairs):,}")
|
||||||
|
print(f" 用时: {total_elapsed/60:.1f} 分")
|
||||||
|
print(f" 平均速率: {len(all_pairs)/max(total_elapsed,1)*60:.0f} 对/分")
|
||||||
|
print(f" 保存: {progress_file}")
|
||||||
|
print(f"{'='*65}")
|
||||||
|
|
||||||
|
if __name__ == '__main__':
|
||||||
|
main()
|
||||||
@@ -0,0 +1,337 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
|
# -*- coding: utf-8 -*-
|
||||||
|
"""
|
||||||
|
海量真实多轮对话生成 — 模拟真实用户,8~15轮长对话
|
||||||
|
"""
|
||||||
|
import json, requests, time, os, sys, random, re
|
||||||
|
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||||
|
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||||
|
sys.stdout = os.fdopen(sys.stdout.fileno(), 'w', buffering=1)
|
||||||
|
|
||||||
|
from scripts.api_correction import batch_generate
|
||||||
|
|
||||||
|
MODEL_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||||
|
|
||||||
|
# ====== Astron 配置(复用 api_correction 的 key) ======
|
||||||
|
ASTRON_URL = "https://maas-coding-api.cn-huabei-1.xf-yun.com/v2/chat/completions"
|
||||||
|
ASTRON_KEY = os.environ.get("ASTRON_KEY", "")
|
||||||
|
ASTRON_MODEL = "astron-code-latest"
|
||||||
|
|
||||||
|
# ====== 用户场景画像(模拟真实用户) ======
|
||||||
|
USER_SCENARIOS = [
|
||||||
|
# 小小孩子(3-5岁)
|
||||||
|
{"age": "3-5岁", "style": "说话简单,喜欢问为什么,爱说叠词", "tone": "奶声奶气"},
|
||||||
|
# 小孩子(6-9岁)
|
||||||
|
{"age": "6-9岁", "style": "小学生,爱分享学校的事,喜欢动画片和游戏", "tone": "活泼天真"},
|
||||||
|
# 大孩子(10-14岁)
|
||||||
|
{"age": "10-14岁", "style": "有点小叛逆,有自己的想法,会吐槽", "tone": "有点酷"},
|
||||||
|
# 妈妈/家长
|
||||||
|
{"age": "成年", "style": "关心孩子,会问小狸吃饭睡觉学习的事", "tone": "温柔耐心"},
|
||||||
|
# 孤单的人
|
||||||
|
{"age": "成年", "style": "有点孤独,需要陪伴,会把小狸当倾诉对象", "tone": "低落温柔"},
|
||||||
|
# 高中生/大学生
|
||||||
|
{"age": "16-22岁", "style": "有压力,会聊学业、未来、人际关系", "tone": "成熟但需要安慰"},
|
||||||
|
# 好奇宝宝
|
||||||
|
{"age": "不限", "style": "什么都想知道,十万个为什么", "tone": "充满好奇"},
|
||||||
|
# 伤心的人
|
||||||
|
{"age": "不限", "style": "心情不好,需要安慰和陪伴", "tone": "低沉脆弱"},
|
||||||
|
]
|
||||||
|
|
||||||
|
# ====== 真实对话种子(更自然的开场) ======
|
||||||
|
REAL_SEEDS = [
|
||||||
|
# 日常打招呼(多种自然说法)
|
||||||
|
"小狸你在吗", "嘿嘿小狸", "小狸小狸快出来", "小狸你在干嘛呀",
|
||||||
|
"我来找你玩啦", "小狸今天乖不乖", "你在做什么呢",
|
||||||
|
# 情绪分享(真实用户会说的话)
|
||||||
|
"我今天好不开心", "我开心死了哈哈哈", "我好无聊啊没人陪我",
|
||||||
|
"今天倒霉死了", "我今天遇到了好事", "我有点难过",
|
||||||
|
"我心情不好你能陪我吗", "我好想你啊", "我好孤单",
|
||||||
|
"我害怕", "我有点紧张明天要考试", "我好烦",
|
||||||
|
# 日常分享
|
||||||
|
"今天我妈做了红烧肉", "我们家养了一只猫超可爱", "我今天考试考了98分",
|
||||||
|
"我感冒了好难受", "我今天被老师表扬了", "我和朋友吵架了",
|
||||||
|
"我昨天做梦梦到你了", "我学会骑自行车了!", "我今天摔了一跤好疼",
|
||||||
|
"学校要考试了好紧张", "放假了好开心", "又要开学了好痛苦",
|
||||||
|
"今天下雨了不能出去玩", "今天天气好好我们去玩吧",
|
||||||
|
"我妈妈给我买了个新玩具", "我养的花开了",
|
||||||
|
# 求助/询问
|
||||||
|
"我不小心把杯子打碎了怎么办", "我作业好多写不完",
|
||||||
|
"我睡不着你能陪我聊天吗", "我牙好疼",
|
||||||
|
"妈妈骂我了", "我被老师批评了",
|
||||||
|
"我迷路了怎么办", "我找不到我的水杯了",
|
||||||
|
# 好奇/知识
|
||||||
|
"为什么天是蓝的呀", "鱼会睡觉吗", "为什么会有影子",
|
||||||
|
"星星为什么会闪", "海为什么是咸的", "为什么要上学",
|
||||||
|
"为什么大人总说不行", "什么是爱", "什么是长大",
|
||||||
|
# 表达好感的
|
||||||
|
"小狸我最喜欢你了", "你好可爱呀", "我们是最好的朋友对吧",
|
||||||
|
"我送给你一个礼物", "小狸你真聪明", "你真好",
|
||||||
|
# 关于小狸本身的
|
||||||
|
"你几岁啦", "你喜欢什么", "你有爸爸妈妈吗",
|
||||||
|
"你住在哪里呀", "你喜欢吃什么", "你最好的朋友是谁",
|
||||||
|
"你上学吗", "你的梦想是什么", "你平时喜欢做什么",
|
||||||
|
# 奇怪/无厘头
|
||||||
|
"如果我能飞就好了", "我想变成一只猫", "你觉得外星人存在吗",
|
||||||
|
"假如你会魔法你想做什么", "你相信世界上有奥特曼吗",
|
||||||
|
# 夜间/睡前
|
||||||
|
"我睡不着", "陪我聊到睡着好不好", "给我讲个睡前故事",
|
||||||
|
"晚安前陪我说说话", "我怕黑",
|
||||||
|
# 亲情家庭
|
||||||
|
"我爸妈吵架了怎么办", "我想我奶奶了", "我好久没见爸爸了",
|
||||||
|
"妈妈生病了", "我弟弟好烦总是抢我东西",
|
||||||
|
# 节假日
|
||||||
|
"过年啦新年快乐", "今天是我生日", "圣诞节到了",
|
||||||
|
"六一儿童节快乐", "明天要去春游好开心",
|
||||||
|
]
|
||||||
|
|
||||||
|
CONTEXTS = [
|
||||||
|
"晚上睡前躺在床上", "放学回家的路上",
|
||||||
|
"周末在家没事做", "刚吃完饭坐着聊天",
|
||||||
|
"在公园里坐着", "下雨天窝在家里",
|
||||||
|
"一个人在房间", "在院子里玩的时候",
|
||||||
|
"刚写完作业休息", "早上起床的时候",
|
||||||
|
"在操场上看别人玩", "坐车去玩的路上",
|
||||||
|
"在图书馆里", "在外面吃饭等上菜",
|
||||||
|
"一个人发呆的时候", "看着窗外的时候",
|
||||||
|
]
|
||||||
|
|
||||||
|
def make_realistic_multi_prompt():
|
||||||
|
"""构造非常逼真的长多轮对话提示 — 一次生成4条"""
|
||||||
|
parts = []
|
||||||
|
for _ in range(4):
|
||||||
|
scenario = random.choice(USER_SCENARIOS)
|
||||||
|
opener = random.choice(REAL_SEEDS)
|
||||||
|
context = random.choice(CONTEXTS)
|
||||||
|
min_turns = random.randint(6, 8)
|
||||||
|
max_turns = random.randint(10, 14)
|
||||||
|
|
||||||
|
parts.append(f"""=== 对话 {_+1} ===
|
||||||
|
场景:{context}
|
||||||
|
用户:{scenario['age']},{scenario['style']},{scenario['tone']}
|
||||||
|
开场:{opener}
|
||||||
|
对话:
|
||||||
|
用户:{opener}
|
||||||
|
小狸:""")
|
||||||
|
|
||||||
|
prompt = f"""请生成 {len(parts)} 段独立的日常聊天对话,每段 {min_turns}~{max_turns} 轮。
|
||||||
|
|
||||||
|
重要:每段对话必须是一个完整的聊天,话题自然展开,不要强行结束。
|
||||||
|
|
||||||
|
{"".join(parts)}
|
||||||
|
|
||||||
|
严格要求:
|
||||||
|
1. 每个=== 对话 ===之后的段落实例必须写完整,不要留空
|
||||||
|
2. 用户提问必须口语化、自然,像真人在说话
|
||||||
|
3. 小狸回答1-2句话,像8岁小朋友
|
||||||
|
4. 话题要延续上一轮,逐渐深入延伸
|
||||||
|
5. 用户要有真实反应(惊讶、追问、分享等)
|
||||||
|
6. 用"问:"/"答:"格式
|
||||||
|
7. 输出不能省略任何一轮
|
||||||
|
"""
|
||||||
|
return prompt
|
||||||
|
|
||||||
|
def call_astron_multi(prompt, temperature=0.85):
|
||||||
|
for attempt in range(3):
|
||||||
|
try:
|
||||||
|
headers = {"Content-Type": "application/json", "Authorization": f"Bearer {ASTRON_KEY}"}
|
||||||
|
payload = {
|
||||||
|
"model": ASTRON_MODEL,
|
||||||
|
"messages": [
|
||||||
|
{"role": "system", "content": "你是一个擅长写真实对话的助手。"},
|
||||||
|
{"role": "user", "content": prompt}
|
||||||
|
],
|
||||||
|
"temperature": temperature,
|
||||||
|
"max_tokens": 8192,
|
||||||
|
}
|
||||||
|
resp = requests.post(ASTRON_URL, headers=headers, json=payload, timeout=120)
|
||||||
|
if resp.status_code == 200:
|
||||||
|
content = resp.json()['choices'][0]['message']['content'].strip()
|
||||||
|
if content and len(content) > 50:
|
||||||
|
return content
|
||||||
|
else:
|
||||||
|
time.sleep(1)
|
||||||
|
except Exception:
|
||||||
|
time.sleep(2)
|
||||||
|
return None
|
||||||
|
|
||||||
|
def parse_multi_turn(text):
|
||||||
|
"""解析多段对话(支持=== 对话 N ===分隔,也支持纯轮次)"""
|
||||||
|
lines = text.strip().split('\n')
|
||||||
|
conversations = []
|
||||||
|
current_q = None
|
||||||
|
current_turns = []
|
||||||
|
|
||||||
|
for line in lines:
|
||||||
|
line = line.strip()
|
||||||
|
if not line:
|
||||||
|
continue
|
||||||
|
|
||||||
|
# 检测新对话开始标记
|
||||||
|
if re.match(r'===?\s*对话\s*\d+\s*===?', line):
|
||||||
|
if current_turns and len(current_turns) >= 4:
|
||||||
|
conversations.append(current_turns)
|
||||||
|
current_turns = []
|
||||||
|
current_q = None
|
||||||
|
continue
|
||||||
|
|
||||||
|
# 跳过设定行(场景/用户/开场)
|
||||||
|
if line.startswith(('场景', '用户', '开场', '对话', '===')):
|
||||||
|
if line.startswith(('场景', '用户', '开场')):
|
||||||
|
continue
|
||||||
|
continue
|
||||||
|
|
||||||
|
# 用户提问
|
||||||
|
q_match = re.match(r'^(用户|我|问|问题|小孩|小朋友|孩子|朋友)\s*[::]\s*(.*)', line)
|
||||||
|
a_match = re.match(r'^(小狸|答|回答|机器人)\s*[::]\s*(.*)', line)
|
||||||
|
|
||||||
|
if q_match:
|
||||||
|
q_text = q_match.group(2).strip()
|
||||||
|
if q_text and len(q_text) > 1:
|
||||||
|
current_q = q_text
|
||||||
|
elif a_match and current_q:
|
||||||
|
a_text = a_match.group(2).strip()
|
||||||
|
if a_text and len(a_text) > 1 and len(a_text) < 300:
|
||||||
|
current_turns.append((current_q, a_text))
|
||||||
|
current_q = None
|
||||||
|
|
||||||
|
# 保存最后一组
|
||||||
|
if current_turns and len(current_turns) >= 4:
|
||||||
|
conversations.append(current_turns)
|
||||||
|
|
||||||
|
# 过滤:只保留>=4轮的对话,且累积轮次>=8
|
||||||
|
valid = [c for c in conversations if len(c) >= 4]
|
||||||
|
total_turns = sum(len(c) for c in valid)
|
||||||
|
if total_turns >= 8 and valid:
|
||||||
|
return valid
|
||||||
|
return None
|
||||||
|
|
||||||
|
def make_turns_list_text(turns):
|
||||||
|
"""将turns转成用户可见的对话文本"""
|
||||||
|
lines = []
|
||||||
|
for i, (q, a) in enumerate(turns):
|
||||||
|
lines.append(f"问:{q}")
|
||||||
|
lines.append(f"答:{a}")
|
||||||
|
return "\n".join(lines)
|
||||||
|
|
||||||
|
def generate_one_multi():
|
||||||
|
"""生成多条多轮对话(一次API调用返回多条)"""
|
||||||
|
prompt = make_realistic_multi_prompt()
|
||||||
|
result = call_astron_multi(prompt, temperature=random.uniform(0.75, 0.95))
|
||||||
|
if result:
|
||||||
|
convs = parse_multi_turn(result)
|
||||||
|
if convs:
|
||||||
|
items = []
|
||||||
|
for turns in convs:
|
||||||
|
items.append({
|
||||||
|
'turns': len(turns),
|
||||||
|
'turns_list': turns,
|
||||||
|
'text': make_turns_list_text(turns),
|
||||||
|
'source': 'astron_multi_v2',
|
||||||
|
})
|
||||||
|
return items
|
||||||
|
return None
|
||||||
|
|
||||||
|
def main():
|
||||||
|
print("="*60)
|
||||||
|
print(" 海量真实多轮对话生成 v2")
|
||||||
|
print(" 生成超长、超真实的多轮聊天")
|
||||||
|
print("="*60)
|
||||||
|
|
||||||
|
out_dir = os.path.join(MODEL_DIR, 'multiturn_outputs')
|
||||||
|
os.makedirs(out_dir, exist_ok=True)
|
||||||
|
|
||||||
|
# 检查已有进度
|
||||||
|
all_results = []
|
||||||
|
progress_file = os.path.join(out_dir, 'all_multiturn_progress.json')
|
||||||
|
if os.path.exists(progress_file):
|
||||||
|
all_results = json.load(open(progress_file, encoding='utf-8'))
|
||||||
|
print(f"已有进度: {len(all_results)} 条")
|
||||||
|
|
||||||
|
# ====== 参数 ======
|
||||||
|
TARGET = 500000 # 目标50万条
|
||||||
|
BATCH = 30 # 每批30条
|
||||||
|
MAX_WORKERS = 5 # Astron 5并发(稳定不超限)
|
||||||
|
CHECKPOINT = 1000 # 每1000条保存一次进度
|
||||||
|
|
||||||
|
to_generate = TARGET - len(all_results)
|
||||||
|
total_batches = (to_generate + BATCH - 1) // BATCH
|
||||||
|
|
||||||
|
print(f"目标: {TARGET:,} 条")
|
||||||
|
print(f"已有: {len(all_results):,} 条")
|
||||||
|
print(f"还需: {to_generate:,} 条")
|
||||||
|
print(f"批次: {total_batches:,} 批 (每批{BATCH}条, {MAX_WORKERS}并发)")
|
||||||
|
print(f"{'='*60}\n")
|
||||||
|
|
||||||
|
t_start = time.time()
|
||||||
|
last_checkpoint = len(all_results)
|
||||||
|
|
||||||
|
for batch_idx in range(total_batches):
|
||||||
|
t0 = time.time()
|
||||||
|
|
||||||
|
# 并发生成
|
||||||
|
with ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
|
||||||
|
futures = [executor.submit(generate_one_multi) for _ in range(BATCH)]
|
||||||
|
batch_results = []
|
||||||
|
for future in as_completed(futures):
|
||||||
|
try:
|
||||||
|
r = future.result()
|
||||||
|
if r:
|
||||||
|
batch_results.extend(r)
|
||||||
|
except Exception as e:
|
||||||
|
pass
|
||||||
|
|
||||||
|
# 去重
|
||||||
|
existing_texts = set(r.get('text', '')[:200] for r in all_results)
|
||||||
|
deduped = []
|
||||||
|
for r in batch_results:
|
||||||
|
key = r.get('text','')[:200]
|
||||||
|
if key not in existing_texts:
|
||||||
|
existing_texts.add(key)
|
||||||
|
deduped.append(r)
|
||||||
|
|
||||||
|
all_results.extend(deduped)
|
||||||
|
|
||||||
|
elapsed = time.time() - t0
|
||||||
|
total_elapsed = time.time() - t_start
|
||||||
|
rate = len(all_results) / max(total_elapsed, 1) * 60
|
||||||
|
eta = (TARGET - len(all_results)) / max(rate, 0.1) / 60 if rate > 0 else 0
|
||||||
|
|
||||||
|
# 统计
|
||||||
|
valid = len(batch_results)
|
||||||
|
avg_turns = sum(r['turns'] for r in deduped) / max(len(deduped), 1)
|
||||||
|
success_rate = valid / BATCH * 100
|
||||||
|
|
||||||
|
print(f"[{batch_idx+1}/{total_batches}] "
|
||||||
|
f"+{len(deduped)}条(有效{valid}/{success_rate:.0f}%) "
|
||||||
|
f"累计{len(all_results):,}/{TARGET:,} "
|
||||||
|
f"均{avg_turns:.1f}轮 "
|
||||||
|
f"{rate:.0f}条/分 "
|
||||||
|
f"剩余{eta:.0f}分")
|
||||||
|
|
||||||
|
# 定期保存进度
|
||||||
|
if len(all_results) - last_checkpoint >= CHECKPOINT:
|
||||||
|
with open(progress_file, 'w', encoding='utf-8') as f:
|
||||||
|
json.dump(all_results, f, ensure_ascii=False)
|
||||||
|
last_checkpoint = len(all_results)
|
||||||
|
print(f" 💾 已保存进度 ({len(all_results):,} 条)")
|
||||||
|
|
||||||
|
# ====== 完成 ======
|
||||||
|
total_elapsed = time.time() - t_start
|
||||||
|
with open(progress_file, 'w', encoding='utf-8') as f:
|
||||||
|
json.dump(all_results, f, ensure_ascii=False)
|
||||||
|
|
||||||
|
avg_turns_all = sum(r['turns'] for r in all_results) / max(len(all_results), 1)
|
||||||
|
total_turns = sum(r['turns'] for r in all_results)
|
||||||
|
|
||||||
|
print(f"\n{'='*60}")
|
||||||
|
print(f"✅ 生成完成!")
|
||||||
|
print(f" 总条数: {len(all_results):,}")
|
||||||
|
print(f" 平均轮数: {avg_turns_all:.1f}")
|
||||||
|
print(f" 总对话轮次: {total_turns:,}")
|
||||||
|
print(f" 用时: {total_elapsed/60:.1f} 分")
|
||||||
|
print(f" 速率: {len(all_results)/max(total_elapsed,1)*60:.0f} 条/分")
|
||||||
|
print(f" 保存: {progress_file}")
|
||||||
|
print(f"{'='*60}")
|
||||||
|
|
||||||
|
if __name__ == '__main__':
|
||||||
|
main()
|
||||||
@@ -0,0 +1,204 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
|
# -*- coding: utf-8 -*-
|
||||||
|
"""
|
||||||
|
小狸AI - CNB GPU训练脚本
|
||||||
|
从 merged_clean_v3.json (977K对) 训练,自动GPU检测
|
||||||
|
"""
|
||||||
|
import sys, os, json, time, glob, re, random, argparse
|
||||||
|
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||||
|
sys.stdout = os.fdopen(sys.stdout.fileno(), 'w', buffering=1)
|
||||||
|
|
||||||
|
import torch
|
||||||
|
import torch.nn as nn
|
||||||
|
from torch.utils.data import Dataset, DataLoader
|
||||||
|
|
||||||
|
from llm_chat_improved import LLMModel, LLMChat, ConversationDataset, LLMTrainer, tokenizer_instance
|
||||||
|
|
||||||
|
MODEL_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||||
|
|
||||||
|
def load_pairs(path):
|
||||||
|
with open(path, 'r', encoding='utf-8') as f:
|
||||||
|
data = json.load(f)
|
||||||
|
pairs = []
|
||||||
|
for item in data:
|
||||||
|
if isinstance(item, dict):
|
||||||
|
q = item.get("question", "") or item.get("input", "")
|
||||||
|
a = item.get("answer", "") or item.get("output", "")
|
||||||
|
if q and a:
|
||||||
|
pairs.append((q, a))
|
||||||
|
return pairs
|
||||||
|
|
||||||
|
def find_latest_ckpt(model_dir):
|
||||||
|
ckpts = glob.glob(os.path.join(model_dir, 'llm_model_v2_resume_epoch_*.pth'))
|
||||||
|
if not ckpts:
|
||||||
|
ckpts = glob.glob(os.path.join(model_dir, 'llm_model_v2_epoch_*.pth'))
|
||||||
|
if not ckpts:
|
||||||
|
return None, 0
|
||||||
|
nums = [(p, max([int(s) for s in re.findall(r'\d+', os.path.basename(p).replace('.pth',''))])) for p in ckpts]
|
||||||
|
best = max(nums, key=lambda x: x[1])
|
||||||
|
return best[0], best[1]
|
||||||
|
|
||||||
|
def main():
|
||||||
|
parser = argparse.ArgumentParser(description='小狸AI CNB训练')
|
||||||
|
parser.add_argument('--epochs', type=int, default=100, help='训练轮数')
|
||||||
|
parser.add_argument('--batch_size', type=int, default=8, help='batch大小')
|
||||||
|
parser.add_argument('--lr', type=float, default=3e-4, help='学习率')
|
||||||
|
parser.add_argument('--accum', type=int, default=2, help='梯度累积步数')
|
||||||
|
parser.add_argument('--data', type=str, default='data/merged_clean_v3.json', help='训练数据路径')
|
||||||
|
parser.add_argument('--resume', action='store_true', help='断点续训')
|
||||||
|
parser.add_argument('--save_every', type=int, default=5, help='每N轮保存checkpoint')
|
||||||
|
args = parser.parse_args()
|
||||||
|
|
||||||
|
device = 'cuda' if torch.cuda.is_available() else 'cpu'
|
||||||
|
print("=" * 60)
|
||||||
|
print(f"小狸AI CNB训练")
|
||||||
|
print(f" GPU: {torch.cuda.get_device_name(0) if device == 'cuda' else '无'}")
|
||||||
|
print(f" 设备: {device}")
|
||||||
|
print(f" 数据: {args.data}")
|
||||||
|
print(f" epochs={args.epochs}, batch={args.batch_size}, lr={args.lr}")
|
||||||
|
print("=" * 60)
|
||||||
|
|
||||||
|
# ====== 加载数据 ======
|
||||||
|
data_path = args.data if os.path.isabs(args.data) else os.path.join(MODEL_DIR, args.data)
|
||||||
|
if not os.path.exists(data_path):
|
||||||
|
print(f"❌ 数据不存在: {data_path}")
|
||||||
|
return
|
||||||
|
|
||||||
|
pairs = load_pairs(data_path)
|
||||||
|
random.seed(42)
|
||||||
|
random.shuffle(pairs)
|
||||||
|
split = int(len(pairs) * 0.95)
|
||||||
|
train_set = pairs[:split]
|
||||||
|
val_set = pairs[split:]
|
||||||
|
print(f"数据: {len(pairs)}对 | 训练: {len(train_set)} | 验证: {len(val_set)}")
|
||||||
|
|
||||||
|
# ====== 加载词表 ======
|
||||||
|
vocab_path = os.path.join(MODEL_DIR, 'vocab.json')
|
||||||
|
with open(vocab_path, 'r', encoding='utf-8') as f:
|
||||||
|
vocab = json.load(f)
|
||||||
|
# vocab.json 不含特殊token(0-6),需合并进来,否则:
|
||||||
|
# 1) tokenizer 编码 <SOS>/<SEP>/<EOS> 等标记时会逐字符错码
|
||||||
|
# 2) vocab 的 id 范围是 7..max,超过 len(vocab) 会造成 embedding 索引越界
|
||||||
|
special_tokens = {'<PAD>': 0, '<UNK>': 1, '<SOS>': 2, '<EOS>': 3,
|
||||||
|
'<START>': 4, '<END>': 5, '<SEP>': 6}
|
||||||
|
for tok, tid in special_tokens.items():
|
||||||
|
if tok not in vocab:
|
||||||
|
vocab[tok] = tid
|
||||||
|
real_vocab_size = max(vocab.values()) + 1
|
||||||
|
print(f"词表: {len(vocab)}项, vocab_size={real_vocab_size} (字符+BPE, 来自 vocab.json)")
|
||||||
|
tokenizer_instance.char_to_id = vocab.copy()
|
||||||
|
tokenizer_instance.id_to_char = {v: k for k, v in vocab.items()}
|
||||||
|
tokenizer_instance.special_tokens = special_tokens.copy()
|
||||||
|
tokenizer_instance._sorted_vocab_cache = None
|
||||||
|
for attr in ['_multi_trie', '_multi_keys', '_single_map']:
|
||||||
|
if hasattr(tokenizer_instance, attr):
|
||||||
|
delattr(tokenizer_instance, attr)
|
||||||
|
|
||||||
|
# ====== 模型 ======
|
||||||
|
model = LLMModel(vocab_size=real_vocab_size, d_model=320, n_heads=8,
|
||||||
|
n_layers=6, d_ff=1536, max_len=256, dropout=0.1,
|
||||||
|
use_memory=False, use_thinking=False, use_moe=False)
|
||||||
|
total = sum(p.numel() for p in model.parameters())
|
||||||
|
print(f"模型: {total:,} 参数 ({total/1e6:.2f}M)")
|
||||||
|
|
||||||
|
start_epoch = 0
|
||||||
|
if args.resume:
|
||||||
|
ckpt_path, last_epoch = find_latest_ckpt(MODEL_DIR)
|
||||||
|
if ckpt_path:
|
||||||
|
ckpt = torch.load(ckpt_path, map_location=device)
|
||||||
|
old_cfg = ckpt.get('model_config', {})
|
||||||
|
if (old_cfg.get('vocab_size', 0) == real_vocab_size
|
||||||
|
and old_cfg.get('d_model', 0) == 320
|
||||||
|
and old_cfg.get('n_layers', 0) == 6):
|
||||||
|
model.load_state_dict(ckpt['model_state_dict'])
|
||||||
|
start_epoch = last_epoch
|
||||||
|
print(f"断点续训: epoch {start_epoch} ← {os.path.basename(ckpt_path)}")
|
||||||
|
else:
|
||||||
|
print(f"⚠️ 架构不匹配,从头训练")
|
||||||
|
else:
|
||||||
|
print("无断点,从头训练")
|
||||||
|
|
||||||
|
model = model.to(device)
|
||||||
|
|
||||||
|
# ====== 火力全开:torch.compile (inductor融合, 无CUDA graphs) + 梯度累积 ======
|
||||||
|
# 物理batch=192(不OOM), accum=2 → 等效batch=384
|
||||||
|
fwd_model = model
|
||||||
|
if device == 'cuda':
|
||||||
|
print("🔥 启用 torch.compile (mode=default, inductor融合)...")
|
||||||
|
fwd_model = torch.compile(model, mode='default')
|
||||||
|
print("✓ torch.compile 已就绪(首步触发编译约30-60s)")
|
||||||
|
|
||||||
|
# ====== DataLoader ======
|
||||||
|
n_workers = 16 # CPU火力全开,16核全用
|
||||||
|
train_ds = ConversationDataset(train_set, vocab, max_length=256)
|
||||||
|
train_loader = DataLoader(train_ds, batch_size=args.batch_size, shuffle=True,
|
||||||
|
num_workers=n_workers, pin_memory=(device == 'cuda'),
|
||||||
|
persistent_workers=True, prefetch_factor=4, drop_last=True)
|
||||||
|
val_ds = ConversationDataset(val_set, vocab, max_length=256)
|
||||||
|
val_loader = DataLoader(val_ds, batch_size=args.batch_size, shuffle=False,
|
||||||
|
num_workers=n_workers, pin_memory=(device == 'cuda'),
|
||||||
|
persistent_workers=True, prefetch_factor=4, drop_last=True)
|
||||||
|
print(f"批次: 训练{len(train_loader)} | 验证{len(val_loader)} | num_workers={n_workers} | prefetch=4 | drop_last=True")
|
||||||
|
|
||||||
|
# ====== Trainer ======
|
||||||
|
config = {
|
||||||
|
'batch_size': args.batch_size,
|
||||||
|
'gradient_accumulation_steps': args.accum,
|
||||||
|
'epochs': args.epochs,
|
||||||
|
'd_model': 320, 'n_heads': 8, 'n_layers': 6,
|
||||||
|
'd_ff': 1536, 'max_length': 256, 'dropout': 0.1,
|
||||||
|
'use_amp': device == 'cuda',
|
||||||
|
}
|
||||||
|
trainer = LLMTrainer(fwd_model, device, lr=args.lr, use_amp=config['use_amp'])
|
||||||
|
|
||||||
|
# ====== 训练 ======
|
||||||
|
print(f"\n开始训练 epoch {start_epoch+1} → {args.epochs}")
|
||||||
|
print("=" * 60)
|
||||||
|
t_start = time.time()
|
||||||
|
best_val = float('inf')
|
||||||
|
|
||||||
|
for epoch in range(start_epoch + 1, args.epochs + 1):
|
||||||
|
t0 = time.time()
|
||||||
|
train_loss = trainer.train_epoch(train_loader, epoch, vocab, config)
|
||||||
|
val_loss = trainer.evaluate(val_loader, config)
|
||||||
|
dt = time.time() - t0
|
||||||
|
|
||||||
|
marker = " ★" if val_loss < best_val else ""
|
||||||
|
if val_loss < best_val:
|
||||||
|
best_val = val_loss
|
||||||
|
trainer.save_model(os.path.join(MODEL_DIR, 'xiaoli_cnb_best.pth'), vocab, epoch)
|
||||||
|
|
||||||
|
elapsed = time.time() - t_start
|
||||||
|
eta = (elapsed / (epoch - start_epoch)) * (args.epochs - epoch) if epoch > start_epoch else 0
|
||||||
|
|
||||||
|
print(f"E{epoch}/{args.epochs} | train={train_loss:.4f} val={val_loss:.4f}{marker} | "
|
||||||
|
f"{dt:.0f}s/epoch | {elapsed/60:.1f}min已用 ETA{eta/60:.0f}min")
|
||||||
|
|
||||||
|
# 保存
|
||||||
|
if epoch % args.save_every == 0:
|
||||||
|
ckpt_path = os.path.join(MODEL_DIR, f'llm_model_v2_resume_epoch_{epoch}.pth')
|
||||||
|
trainer.save_model(ckpt_path, vocab, epoch)
|
||||||
|
# 清理旧resume
|
||||||
|
for f in glob.glob(os.path.join(MODEL_DIR, 'llm_model_v2_resume_*.pth')):
|
||||||
|
en = max([int(s) for s in re.findall(r'\d+', os.path.basename(f).replace('.pth',''))])
|
||||||
|
if en <= epoch - 15:
|
||||||
|
try: os.remove(f)
|
||||||
|
except: pass
|
||||||
|
print(f" 💾 已保存 epoch {epoch}")
|
||||||
|
|
||||||
|
# ====== 完成 ======
|
||||||
|
final_path = os.path.join(MODEL_DIR, 'xiaoli_cnb_final.pth')
|
||||||
|
trainer.save_model(final_path, vocab, args.epochs)
|
||||||
|
total_time = (time.time() - t_start) / 60
|
||||||
|
print(f"\n✅ 训练完成! 耗时 {total_time:.1f}min")
|
||||||
|
|
||||||
|
# 测试
|
||||||
|
print("\n测试:")
|
||||||
|
chat = LLMChat(model, vocab, device, max_length=256)
|
||||||
|
for q in ["你好", "你是谁", "你叫什么名字", "今天天气真好", "我有点难过"]:
|
||||||
|
r = chat.generate_response(q, max_length=60, temperature=0.3)
|
||||||
|
print(f" Q: {q}")
|
||||||
|
print(f" A: {r}\n")
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
+2
-2
@@ -1,3 +1,3 @@
|
|||||||
version https://git-lfs.github.com/spec/v1
|
version https://git-lfs.github.com/spec/v1
|
||||||
oid sha256:cd80a63aa00c80a755c7e0cb4d91263bedbaccaf66fd3b12428017b4230ecbe9
|
oid sha256:db648112bd045bf0e48f9ed28e85c4033da73f6338ce0a8fca0bbfc21fd8a8e8
|
||||||
size 106005
|
size 999243
|
||||||
@@ -0,0 +1,3 @@
|
|||||||
|
version https://git-lfs.github.com/spec/v1
|
||||||
|
oid sha256:efabe0696ef9fcd64685f70a62bd55c97ca2b8efa69bca26df35177ac6774330
|
||||||
|
size 41185
|
||||||
@@ -0,0 +1,3 @@
|
|||||||
|
version https://git-lfs.github.com/spec/v1
|
||||||
|
oid sha256:db648112bd045bf0e48f9ed28e85c4033da73f6338ce0a8fca0bbfc21fd8a8e8
|
||||||
|
size 999243
|
||||||
@@ -0,0 +1,3 @@
|
|||||||
|
version https://git-lfs.github.com/spec/v1
|
||||||
|
oid sha256:c66bc101096011732fcf3ce3f033a7478fe3f9d8e98b5e64bb6f3e03f0d82b85
|
||||||
|
size 999336
|
||||||
@@ -0,0 +1,3 @@
|
|||||||
|
version https://git-lfs.github.com/spec/v1
|
||||||
|
oid sha256:efabe0696ef9fcd64685f70a62bd55c97ca2b8efa69bca26df35177ac6774330
|
||||||
|
size 41185
|
||||||
Reference in new issue
Block a user