# 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", }