Upload hermes_core/trainer.py with huggingface_hub

hermes_core/trainer.py

@@ -0,0 +1,280 @@
+"""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 ""