guanghulab/zhuyuan-agent/training_runner.py

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