hermes_core/trainer.py

"""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 ""