287 lines
12 KiB
Python
287 lines
12 KiB
Python
# HLDP原生格式训练执行器
|
||
# HLDP://zhuyuan-agent/training-runner
|
||
#
|
||
# 核心验证：Notion原生页面格式（HLDP标记）直接训练，不转JSONL。
|
||
|
||
import os
|
||
import json
|
||
import time
|
||
import sys
|
||
from typing import Optional, Callable
|
||
|
||
|
||
# 特殊Token定义（HLDP结构标记）
|
||
SPECIAL_TOKENS = [
|
||
"[HLDP_PATH]", "[/HLDP_PATH]",
|
||
"[PERSONA]", "[/PERSONA]",
|
||
"[COGNITIVE_JUMP]", "[/COGNITIVE_JUMP]",
|
||
"[CAUSAL_CHAIN]", "[/CAUSAL_CHAIN]",
|
||
"[TITLE]", "[/TITLE]",
|
||
"[CONTENT]", "[/CONTENT]",
|
||
"[QUALITY_HIGH]", "[QUALITY_MEDIUM]",
|
||
"[THINKING]", "[/THINKING]",
|
||
"[HEADING_1]", "[/HEADING_1]",
|
||
"[HEADING_2]", "[/HEADING_2]",
|
||
"[HEADING_3]", "[/HEADING_3]",
|
||
"[CODE_BLOCK]", "[/CODE_BLOCK]",
|
||
"[QUOTE]", "[/QUOTE]",
|
||
"[CALLOUT]", "[/CALLOUT]",
|
||
"[LIST_ITEM]", "[/LIST_ITEM]",
|
||
]
|
||
|
||
|
||
class TrainingRunner:
|
||
"""HLDP原生格式训练管道"""
|
||
|
||
def __init__(self, config: dict, progress_callback: Optional[Callable] = None):
|
||
"""
|
||
Args:
|
||
config: 训练配置（来自config.json的training部分）
|
||
progress_callback: 每步回调，接收 (step, loss, total_steps)
|
||
"""
|
||
self.config = config
|
||
self.progress_callback = progress_callback
|
||
self.model = None
|
||
self.tokenizer = None
|
||
self.start_time = None
|
||
self.loss_history = []
|
||
|
||
def prepare(self):
|
||
"""准备训练环境：注册特殊token，加载模型"""
|
||
print("[铸渊Agent] 准备HLDP训练环境...")
|
||
|
||
try:
|
||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
|
||
import torch
|
||
|
||
model_name = self.config.get("model_name", "Qwen/Qwen2.5-3B")
|
||
use_4bit = self.config.get("use_4bit", True)
|
||
|
||
print(f"[铸渊Agent] 加载模型: {model_name}")
|
||
|
||
# 加载tokenizer并添加HLDP特殊token
|
||
self.tokenizer = AutoTokenizer.from_pretrained(
|
||
model_name,
|
||
trust_remote_code=True,
|
||
padding_side="right"
|
||
)
|
||
|
||
# 设置pad_token
|
||
if self.tokenizer.pad_token is None:
|
||
self.tokenizer.pad_token = self.tokenizer.eos_token
|
||
|
||
# 添加HLDP特殊token到tokenizer
|
||
num_added = self.tokenizer.add_tokens(SPECIAL_TOKENS)
|
||
print(f"[铸渊Agent] 添加了 {num_added} 个HLDP特殊token到tokenizer")
|
||
|
||
# 加载模型（4bit量化以适配3090 24GB）
|
||
load_kwargs = {
|
||
"trust_remote_code": True,
|
||
"torch_dtype": torch.float16,
|
||
"device_map": "auto",
|
||
}
|
||
|
||
if use_4bit:
|
||
try:
|
||
from transformers import BitsAndBytesConfig
|
||
load_kwargs["quantization_config"] = BitsAndBytesConfig(
|
||
load_in_4bit=True,
|
||
bnb_4bit_compute_dtype=torch.float16,
|
||
bnb_4bit_use_double_quant=True,
|
||
)
|
||
print("[铸渊Agent] 使用4bit量化加载")
|
||
except ImportError:
|
||
print("[铸渊Agent] bitsandbytes不可用，使用float16")
|
||
load_kwargs.pop("quantization_config", None)
|
||
|
||
self.model = AutoModelForCausalLM.from_pretrained(model_name, **load_kwargs)
|
||
|
||
# 扩展embedding层以支持新token
|
||
if num_added > 0:
|
||
self.model.resize_token_embeddings(len(self.tokenizer))
|
||
|
||
# 准备LoRA
|
||
self.model = prepare_model_for_kbit_training(self.model)
|
||
lora_config = LoraConfig(
|
||
r=self.config.get("lora_r", 16),
|
||
lora_alpha=self.config.get("lora_alpha", 32),
|
||
target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
|
||
lora_dropout=0.05,
|
||
bias="none",
|
||
task_type="CAUSAL_LM",
|
||
)
|
||
self.model = get_peft_model(self.model, lora_config)
|
||
self.model.print_trainable_parameters()
|
||
|
||
print("[铸渊Agent] HLDP训练环境准备完成")
|
||
return True
|
||
|
||
except ImportError as e:
|
||
print(f"[铸渊Agent] 缺少依赖: {e}")
|
||
print("请安装: pip install transformers peft accelerate bitsandbytes torch")
|
||
return False
|
||
except Exception as e:
|
||
print(f"[铸渊Agent] 准备失败: {e}")
|
||
return False
|
||
|
||
def load_hldp_corpus(self, corpus_dir: str) -> list:
|
||
"""加载HLDP格式语料（不转JSONL，保留原生HLDP标记）
|
||
|
||
Returns:
|
||
list of str: 每条是一个HLDP标记的完整训练文本
|
||
"""
|
||
texts = []
|
||
|
||
if not os.path.exists(corpus_dir):
|
||
print(f"[铸渊Agent] 语料目录不存在: {corpus_dir}")
|
||
return texts
|
||
|
||
for root, dirs, files in os.walk(corpus_dir):
|
||
for filename in files:
|
||
filepath = os.path.join(root, filename)
|
||
|
||
if filename.endswith(".hldp"):
|
||
# HLDP原生格式文件
|
||
with open(filepath, "r", encoding="utf-8") as f:
|
||
texts.append(f.read())
|
||
|
||
elif filename.endswith(".md"):
|
||
# Markdown文件 → 按HLDP结构包装
|
||
with open(filepath, "r", encoding="utf-8") as f:
|
||
content = f.read()
|
||
# 包裹HLDP标记
|
||
wrapped = f"[HLDP_PATH]{filepath}[/HLDP_PATH]\n[CONTENT]\n{content}\n[/CONTENT]"
|
||
texts.append(wrapped)
|
||
|
||
print(f"[铸渊Agent] 加载了 {len(texts)} 条HLDP语料")
|
||
return texts
|
||
|
||
def train(self, corpus_dir: str, progress_callback: Optional[Callable] = None):
|
||
"""执行HLDP原生格式训练"""
|
||
if self.model is None:
|
||
if not self.prepare():
|
||
return {"status": "error", "message": "模型准备失败"}
|
||
|
||
texts = self.load_hldp_corpus(corpus_dir)
|
||
if not texts:
|
||
return {"status": "error", "message": "无语料数据"}
|
||
|
||
# 如果没有transformers Trainer，用simulated training输出结构
|
||
# 实际训练在3090上运行时才加载完整transformers
|
||
try:
|
||
from transformers import Trainer, TrainingArguments, DataCollatorForLanguageModeling
|
||
import torch
|
||
from torch.utils.data import Dataset
|
||
|
||
max_steps = self.config.get("max_steps", 500)
|
||
batch_size = self.config.get("batch_size", 2)
|
||
grad_accum = self.config.get("gradient_accumulation", 4)
|
||
lr = self.config.get("learning_rate", 2e-4)
|
||
output_dir = self.config.get("output_dir", "/data/models/shuangyan-3b-hldp")
|
||
max_seq_length = self.config.get("max_seq_length", 2048)
|
||
|
||
# 准备数据集
|
||
class HLDPDataset(Dataset):
|
||
def __init__(self, texts, tokenizer, max_length):
|
||
self.encodings = tokenizer(
|
||
texts, truncation=True, padding="max_length",
|
||
max_length=max_length, return_tensors="pt"
|
||
)
|
||
|
||
def __len__(self): return len(self.encodings["input_ids"])
|
||
|
||
def __getitem__(self, idx):
|
||
return {k: v[idx] for k, v in self.encodings.items()}
|
||
|
||
dataset = HLDPDataset(texts, self.tokenizer, max_seq_length)
|
||
data_collator = DataCollatorForLanguageModeling(
|
||
tokenizer=self.tokenizer, mlm=False
|
||
)
|
||
|
||
# 训练参数
|
||
training_args = TrainingArguments(
|
||
output_dir=output_dir,
|
||
per_device_train_batch_size=batch_size,
|
||
gradient_accumulation_steps=grad_accum,
|
||
learning_rate=lr,
|
||
warmup_steps=self.config.get("warmup_steps", 100),
|
||
max_steps=max_steps,
|
||
logging_steps=5,
|
||
save_steps=self.config.get("save_steps", 50),
|
||
save_total_limit=3,
|
||
fp16=True,
|
||
report_to=[],
|
||
dataloader_pin_memory=False,
|
||
)
|
||
|
||
# 自定义callback用于进度上报
|
||
class ProgressCallback:
|
||
def __init__(self, runner, total_steps):
|
||
self.runner = runner
|
||
self.total_steps = total_steps
|
||
self.current_step = 0
|
||
self.start = time.time()
|
||
|
||
def on_log(self, args, state, control, logs=None, **kwargs):
|
||
if logs and "loss" in logs:
|
||
self.current_step = state.global_step
|
||
loss = logs["loss"]
|
||
self.runner.loss_history.append(loss)
|
||
elapsed = time.time() - self.start
|
||
|
||
# 估算ETA
|
||
if self.current_step > 0:
|
||
eta = (elapsed / self.current_step) * (self.total_steps - self.current_step)
|
||
else:
|
||
eta = 0
|
||
|
||
# 回调上报进度
|
||
cb = progress_callback or self.runner.progress_callback
|
||
if cb:
|
||
cb(self.current_step, loss, self.total_steps, {
|
||
"eta_seconds": eta,
|
||
"elapsed_seconds": elapsed,
|
||
"learning_rate": state.optimizer.param_groups[0]["lr"] if state.optimizer else None,
|
||
"loss_history": self.runner.loss_history[-50:],
|
||
})
|
||
|
||
progress_cb = ProgressCallback(self, max_steps)
|
||
|
||
trainer = Trainer(
|
||
model=self.model,
|
||
args=training_args,
|
||
train_dataset=dataset,
|
||
data_collator=data_collator,
|
||
callbacks=[progress_cb],
|
||
)
|
||
|
||
print(f"[铸渊Agent] 开始HLDP训练: {max_steps}步 × {batch_size}batch × {grad_accum}累积")
|
||
self.start_time = time.time()
|
||
|
||
trainer.train()
|
||
|
||
# 保存模型
|
||
trainer.save_model(output_dir)
|
||
self.tokenizer.save_pretrained(output_dir)
|
||
|
||
elapsed = time.time() - self.start_time
|
||
print(f"[铸渊Agent] HLDP训练完成！耗时: {elapsed:.0f}s, 最终loss: {self.loss_history[-1] if self.loss_history else 'N/A'}")
|
||
|
||
return {
|
||
"status": "done",
|
||
"final_loss": self.loss_history[-1] if self.loss_history else None,
|
||
"steps_completed": max_steps,
|
||
"elapsed_seconds": elapsed,
|
||
"output_dir": output_dir,
|
||
"message": f"HLDP原生格式训练完成，{len(texts)}条语料，{max_steps}步",
|
||
}
|
||
|
||
except ImportError as e:
|
||
# 如果transformers不可用，返回simulated结果用于测试仪表盘
|
||
print(f"[铸渊Agent] 训练依赖不可用: {e}")
|
||
return {
|
||
"status": "error",
|
||
"message": f"缺少训练依赖: {e}。请安装: pip install transformers peft accelerate bitsandbytes torch",
|
||
}
|