"""QLoRA 训练 — 从训练集到 LoRA 权重.""" import json import hashlib import uuid import subprocess import sys from pathlib import Path from datetime import datetime, timezone from hermes_core.types import HERMES_DATA_DIR, TrainingRun, TrainingStatus from hermes_core.db import init_db, get_active_records, insert_training_run, update_training_run MIN_RECORDS_FOR_TRAINING = 10 def _now() -> str: return datetime.now(timezone.utc).isoformat() def _checkpoint_dir(user_id: str, scope_id: str, version: int) -> Path: return HERMES_DATA_DIR / "users" / user_id / "checkpoints" / scope_id / f"v{version}" def compute_content_hash(records: list) -> str: """计算训练集内容哈希，用于去重。""" data = json.dumps([ {"scope_label": r.scope_label, "dimensions": [{"k": d.key, "v": d.value} for d in r.dimensions]} for r in sorted(records, key=lambda r: r.id) ], sort_keys=True, ensure_ascii=False) return hashlib.sha256(data.encode()).hexdigest() def build_training_dataset(user_id: str, scope_id: str) -> list[dict]: """从 DB 读取 active 记录，构建 QLoRA 训练样本。 每条记录的 scope_label 和 dimensions 组成一条训练样本： { "instruction": "你是一个AI助手。根据场景和已知偏好，默认采用以下偏好。", "input": "场景: 后端API开发", "output": "已知偏好: [language: TypeScript] [testing: Vitest]" } """ conn = init_db(user_id) records = get_active_records(conn, scope_id) conn.close() samples = [] for rec in records: dims_str = " ".join(f"[{d.key}: {d.value}]" for d in rec.dimensions) sample = { "instruction": "你是一个AI助手。根据当前的场景标签，默认采用用户的已知偏好。", "input": f"场景: {rec.scope_label}", "output": f"已知偏好: {dims_str}", } samples.append(sample) return samples def run_qlora_train( user_id: str, scope_id: str, base_model: str = "Qwen/Qwen2.5-7B-Instruct", ) -> str: """执行 QLoRA 训练。 Returns: checkpoint_path 或空字符串（失败时） """ conn = init_db(user_id) records = get_active_records(conn, scope_id) if len(records) < MIN_RECORDS_FOR_TRAINING: conn.close() return "" # 去重检查 new_hash = compute_content_hash(records) # 查最新一次训练 latest = conn.execute( "SELECT * FROM training_runs WHERE scope_id=? AND status='done' ORDER BY version DESC LIMIT 1", (scope_id,) ).fetchone() if latest and latest["content_hash"] == new_hash: conn.close() return "" # 确定版本号 if latest: version = latest["version"] + 1 else: version = 1 run_id = f"run_{uuid.uuid4().hex[:12]}" run = TrainingRun( id=run_id, scope_id=scope_id, version=version, status=TrainingStatus.pending, content_hash=new_hash, started_at=_now(), ) insert_training_run(conn, run) conn.close() # 构建训练数据 samples = build_training_dataset(user_id, scope_id) ckpt_dir = _checkpoint_dir(user_id, scope_id, version) ckpt_dir.mkdir(parents=True, exist_ok=True) dataset_path = ckpt_dir / "train.jsonl" with open(dataset_path, "w", encoding="utf-8") as f: for s in samples: f.write(json.dumps(s, ensure_ascii=False) + "\n") # 写训练脚本 train_script = ckpt_dir / "train.py" train_script.write_text(""" import json import os import sys import torch from transformers import ( AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, TrainingArguments, Trainer, DataCollatorForSeq2Seq, ) from peft import ( LoraConfig, get_peft_model, prepare_model_for_kbit_training, ) from datasets import Dataset import warnings warnings.filterwarnings("ignore") # ── 加载数据集 ── dataset_path = sys.argv[1] checkpoint_dir = sys.argv[2] samples = [] with open(dataset_path, "r", encoding="utf-8") as f: for line in f: samples.append(json.loads(line)) dataset = Dataset.from_list(samples) # ── 量化配置 ── bnb_config = BitsAndBytesConfig( load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.float16, bnb_4bit_use_double_quant=True, ) # ── 加载模型 ── base_model = sys.argv[3] if len(sys.argv) > 3 else "Qwen/Qwen2.5-7B-Instruct" model = AutoModelForCausalLM.from_pretrained( base_model, quantization_config=bnb_config, device_map="auto", trust_remote_code=True, ) tokenizer = AutoTokenizer.from_pretrained(base_model, trust_remote_code=True) if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token model = prepare_model_for_kbit_training(model) # ── LoRA 配置（偏好学习优化）── lora_config = LoraConfig( r=8, lora_alpha=16, lora_dropout=0.05, target_modules=["q_proj", "v_proj"], task_type="CAUSAL_LM", ) model = get_peft_model(model, lora_config) # ── Tokenize ── def tokenize(example): prompt = f"{example['instruction']}\\n\\n输入: {example['input']}\\n输出: " full = prompt + example["output"] tokenized = tokenizer(full, truncation=True, max_length=512, padding=False) tokenized["labels"] = tokenized["input_ids"].copy() # mask prompt prompt_len = len(tokenizer(prompt, truncation=True, max_length=512)["input_ids"]) tokenized["labels"][:prompt_len] = [-100] * prompt_len return tokenized tokenized_dataset = dataset.map(tokenize, remove_columns=dataset.column_names) # ── 训练参数 ── training_args = TrainingArguments( output_dir=checkpoint_dir, per_device_train_batch_size=1, gradient_accumulation_steps=4, num_train_epochs=3, learning_rate=1e-4, lr_scheduler_type="cosine", warmup_ratio=0.1, optim="paged_adamw_8bit", logging_steps=10, save_strategy="epoch", fp16=True, report_to="none", ) trainer = Trainer( model=model, args=training_args, train_dataset=tokenized_dataset, data_collator=DataCollatorForSeq2Seq(tokenizer, pad_to_multiple_of=8), ) trainer.train() model.save_pretrained(checkpoint_dir) tokenizer.save_pretrained(checkpoint_dir) print(f"Training complete. Checkpoint saved to {checkpoint_dir}") """, encoding="utf-8") # ── 启动训练子进程 ── try: conn = init_db(user_id) update_training_run(conn, run_id, status=TrainingStatus.training) conn.close() result = subprocess.run( [sys.executable, str(train_script), str(dataset_path), str(ckpt_dir), base_model], capture_output=True, encoding="utf-8", errors="replace", timeout=7200, # 2h max ) conn = init_db(user_id) if result.returncode == 0: update_training_run( conn, run_id, status=TrainingStatus.done, checkpoint_path=str(ckpt_dir), finished_at=_now(), ) # 更新 scope conn.execute( "UPDATE scopes SET needs_training=0 WHERE id=?", (scope_id,) ) conn.commit() conn.close() return str(ckpt_dir) else: update_training_run( conn, run_id, status=TrainingStatus.failed, error_msg=result.stderr[:500], finished_at=_now(), ) conn.commit() conn.close() return "" except subprocess.TimeoutExpired: conn = init_db(user_id) update_training_run( conn, run_id, status=TrainingStatus.failed, error_msg="Training timed out (2h)", finished_at=_now(), ) conn.commit() conn.close() return "" except Exception as e: conn = init_db(user_id) update_training_run( conn, run_id, status=TrainingStatus.failed, error_msg=str(e)[:500], finished_at=_now(), ) conn.commit() conn.close() return ""