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()