ai-agent-book 精选快照（<2MB 代码与文档，来自 github.com/bojieli/ai-agent-book）

chapter8/exact-copy-sft/train_sft.py

@@ -0,0 +1,60 @@
+from __future__ import annotations
+import argparse, hashlib, json
+from datetime import datetime, timezone
+from pathlib import Path
+import torch
+from peft import LoraConfig, get_peft_model
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+ROOT = Path(__file__).resolve().parent
+
+def load(path):
+    return [json.loads(x) for x in Path(path).read_text(encoding="utf-8").splitlines() if x.strip()]
+
+def encode(tokenizer, row, max_len):
+    messages = [
+        {"role": "system", "content": "你是工具调用中的精确复制器。任何不透明字符串都必须逐字保留。"},
+        {"role": "user", "content": row["prompt"]},
+    ]
+    prompt_ids = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, enable_thinking=False)
+    target_ids = tokenizer(row["target"], add_special_tokens=False).input_ids
+    eos = [] if tokenizer.eos_token_id is None else [tokenizer.eos_token_id]
+    ids = (prompt_ids + target_ids + eos)[:max_len]
+    return {"input_ids": ids, "labels": ([-100] * len(prompt_ids) + target_ids + eos)[:max_len], "attention_mask": [1] * len(ids)}
+
+class Collator:
+    def __init__(self, pad): self.pad = pad
+    def __call__(self, batch):
+        m = max(len(x["input_ids"]) for x in batch)
+        return {
+            "input_ids": torch.tensor([x["input_ids"] + [self.pad] * (m - len(x["input_ids"])) for x in batch]),
+            "labels": torch.tensor([x["labels"] + [-100] * (m - len(x["labels"])) for x in batch]),
+            "attention_mask": torch.tensor([x["attention_mask"] + [0] * (m - len(x["attention_mask"])) for x in batch]),
+        }
+
+def main():
+    ap = argparse.ArgumentParser(); ap.add_argument("--model", default="Qwen/Qwen3-8B"); ap.add_argument("--epochs", type=int, default=2); ap.add_argument("--lr", type=float, default=1e-4); ap.add_argument("--batch-size", type=int, default=2); ap.add_argument("--grad-accum", type=int, default=4); ap.add_argument("--max-length", type=int, default=768); ap.add_argument("--seed", type=int, default=719); ap.add_argument("--output", default=str(ROOT / "output" / "adapter")); args = ap.parse_args()
+    if not torch.cuda.is_available(): raise SystemExit("需要 CUDA GPU")
+    torch.manual_seed(args.seed)
+    data_path = ROOT / "data" / "train.jsonl"; rows = load(data_path)
+    tokenizer = AutoTokenizer.from_pretrained(args.model, use_fast=True); tokenizer.pad_token = tokenizer.pad_token or tokenizer.eos_token
+    encoded = [encode(tokenizer, r, args.max_length) for r in rows]
+    model = AutoModelForCausalLM.from_pretrained(args.model, torch_dtype=torch.bfloat16, device_map="auto"); model.config.use_cache = False
+    model = get_peft_model(model, LoraConfig(r=16, lora_alpha=32, lora_dropout=0.05, bias="none", task_type="CAUSAL_LM", target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]))
+    model.print_trainable_parameters()
+    loader = torch.utils.data.DataLoader(encoded, batch_size=args.batch_size, shuffle=True, collate_fn=Collator(tokenizer.pad_token_id))
+    opt = torch.optim.AdamW((p for p in model.parameters() if p.requires_grad), lr=args.lr); device = next(p for p in model.parameters() if p.requires_grad).device
+    model.train(); losses = []; started = datetime.now(timezone.utc)
+    for epoch in range(args.epochs):
+        opt.zero_grad(set_to_none=True)
+        for step, batch in enumerate(loader):
+            batch = {k: v.to(device) for k, v in batch.items()}; loss = model(**batch).loss / args.grad_accum; loss.backward()
+            if (step + 1) % args.grad_accum == 0 or step == len(loader) - 1:
+                torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0); opt.step(); opt.zero_grad(set_to_none=True); losses.append(float(loss.detach().cpu() * args.grad_accum))
+        print(f"epoch={epoch + 1}/{args.epochs} loss={losses[-1]:.4f}", flush=True)
+    out = Path(args.output); out.mkdir(parents=True, exist_ok=True); model.save_pretrained(out); tokenizer.save_pretrained(out)
+    finished = datetime.now(timezone.utc); run = started.strftime("train_%Y%m%dT%H%M%SZ"); vdir = ROOT / "validation" / run; vdir.mkdir(parents=True, exist_ok=True)
+    receipt = {"experiment": "8-19-exact-copy-sft", "run": run, "model": args.model, "data_sha256": hashlib.sha256(data_path.read_bytes()).hexdigest(), "train_examples": len(rows), "config": vars(args), "cuda": torch.cuda.get_device_name(0), "started_at": started.isoformat(), "finished_at": finished.isoformat(), "final_loss": losses[-1]}
+    (vdir / "training_receipt.json").write_text(json.dumps(receipt, ensure_ascii=False, indent=2), encoding="utf-8"); (ROOT / "validation" / "latest.json").write_text(json.dumps({"run": run, "training_receipt": str((vdir / "training_receipt.json").relative_to(ROOT))}, ensure_ascii=False, indent=2), encoding="utf-8"); print(json.dumps(receipt, ensure_ascii=False, indent=2))
+
+if __name__ == "__main__": main()