scripts/train_decision_from_lm.py
Rama main · réplica de lectura del repositorio original.
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved.
"""Turn a plain language model into a Jev-style decision model with Unsloth.
python scripts/train_decision_from_lm.py --model unsloth/Qwen3.5-2B --rows 12000 \
--out outputs/qwen35_2b_decision
Recipe (see FastDecisionModel docs): a fresh Clef joint schema head on the backbone's last
hidden states, an optional head-only warm-up with the backbone frozen, then LoRA on attention,
MLP and Gated DeltaNet projections plus the head, soft-target cross-entropy, and a temperature
fitted on held-out decisions. Prints one JSON line with the results.
"""
import argparse
import json
import random
import time
from unsloth import DecisionTrainer, FastDecisionModel # noqa: I001 (Unsloth first)
from unsloth.models import decision_datasets as dd
import torch
from transformers import TrainingArguments
DEFAULT_SOURCES = (
"banking77",
"clinc150",
"mnli",
"snli",
"wanli",
"boolq",
"ag_news",
"sst5",
"mmlu",
"commonsense_qa",
"arc",
"prompt_injections",
"xlam",
)
def _record_metrics(model, tokenizer, items) -> dict:
metrics = FastDecisionModel.evaluate(model, tokenizer, items)
return {
"decisions": sum(len(item.get("labels", [item.get("label")])) for item in items),
**{k: round(float(metrics[k]), 4) for k in ("accuracy", "ece", "record_accuracy")},
}
def main():
parser = argparse.ArgumentParser(description = __doc__)
parser.add_argument("--model", default = "unsloth/Qwen3.5-2B")
parser.add_argument(
"--sources", default = ",".join(DEFAULT_SOURCES), help = "comma separated, or 'none'"
)
parser.add_argument(
"--rows", type = int, default = 12000, help = "mixture rows besides typed-decisions"
)
parser.add_argument(
"--typed-decisions", type = int, default = 1, help = "include LocalLLaMA/typed-decisions train"
)
parser.add_argument("--epochs", type = float, default = 1.0)
parser.add_argument("--head-warmup-steps", type = int, default = 0)
parser.add_argument("--head-width", type = int, default = None)
parser.add_argument("--lora-r", type = int, default = 64)
parser.add_argument("--lr", type = float, default = 1e-4)
parser.add_argument("--head-lr", type = float, default = 3e-4)
parser.add_argument("--batch-size", type = int, default = 8)
parser.add_argument("--grad-accum", type = int, default = 2)
parser.add_argument("--max-seq-length", type = int, default = 4096)
parser.add_argument("--load-in-4bit", action = "store_true")
parser.add_argument("--eval-rows", type = int, default = 500, help = "BANKING77 / CLINC150 test rows")
parser.add_argument("--seed", type = int, default = 3407)
parser.add_argument("--out", default = None, help = "save the merged decision model here")
parser.add_argument("--json", default = None, help = "also write the result JSON here")
args = parser.parse_args()
start = time.time()
model, processor = FastDecisionModel.from_pretrained(
args.model,
decision_head = "clef",
head_width = args.head_width,
max_seq_length = args.max_seq_length,
load_in_4bit = args.load_in_4bit,
random_state = args.seed,
)
# Evaluation sets first, so the training mixture can be decontaminated against them.
td_test = dd.load_source("typed_decisions", "test")
evals = {"typed_decisions": td_test}
for name in ("banking77", "clinc150"):
evals[name] = [
dd.canonical_row(r) for r in dd.load_source(name, "test", limit = args.eval_rows, seed = 0)
]
# Raw states: the decontaminator serializes both sides with sorted keys.
eval_texts = [r["state"] for rows in evals.values() for r in rows]
pools = {}
sources = [] if args.sources == "none" else args.sources.split(",")
for name in sources:
try:
# Gated sources fail first, so the survivors' share is known before the full load.
pools[name] = dd.load_source(name, "train", limit = 1, seed = args.seed)
except Exception as exc: # gated or unavailable sources are reported and skipped
print(f"skipping {name}: {type(exc).__name__}: {str(exc)[:160]}")
per_source = max(1, -(-args.rows // max(1, len(pools))))
# 2x headroom so decontamination and converter drops do not force resampling.
pools = {
name: dd.load_source(name, "train", limit = 2 * per_source, seed = args.seed) for name in pools
}
rows = (
dd.build_decision_mixture(
pools, n_rows = args.rows, seed = args.seed, decontaminate_against = eval_texts
)
if pools
else []
)
if args.typed_decisions:
rng = random.Random(args.seed)
cleaner = dd.Decontaminator(eval_texts)
rows += [
dd.augment_row(r, rng)
for r in dd.load_source("typed_decisions", "train")
if not cleaner.contaminated(r)
]
random.Random(args.seed).shuffle(rows)
items, report = FastDecisionModel.build_dataset(rows, processor, model)
train, holdout = FastDecisionModel.split_holdout(items, args.seed)
eval_items = {}
for name, eval_rows in evals.items():
eval_items[name], eval_report = FastDecisionModel.build_dataset(eval_rows, processor, model)
base = {name: _record_metrics(model, processor, its) for name, its in eval_items.items()}
def trainer(steps = None, epochs = None):
return DecisionTrainer(
model = model,
tokenizer = processor,
train_dataset = train,
head_learning_rate = args.head_lr,
args = TrainingArguments(
output_dir = "outputs/_decision_from_lm_run",
per_device_train_batch_size = args.batch_size,
gradient_accumulation_steps = args.grad_accum,
learning_rate = args.lr,
lr_scheduler_type = "cosine",
warmup_steps = max(
1,
round(
0.05
* (steps or epochs * -(-len(train) // (args.batch_size * args.grad_accum)))
),
),
weight_decay = 0.01,
max_steps = steps or -1,
num_train_epochs = epochs or 1,
bf16 = torch.cuda.is_available() and torch.cuda.is_bf16_supported(),
logging_steps = 10,
report_to = "none",
save_strategy = "no",
seed = args.seed,
),
)
if torch.cuda.is_available():
torch.cuda.reset_peak_memory_stats()
warmup = None
if args.head_warmup_steps:
FastDecisionModel.freeze_backbone(model)
out = trainer(steps = args.head_warmup_steps).train()
warmup = {
"steps": args.head_warmup_steps,
"s_per_step": round(out.metrics["train_runtime"] / args.head_warmup_steps, 3),
}
FastDecisionModel.unfreeze_backbone(model)
model = FastDecisionModel.get_peft_model(
model, r = args.lora_r, lora_alpha = args.lora_r, random_state = args.seed
)
t = trainer(epochs = args.epochs)
out = t.train()
steps = t.state.global_step
losses = [h["loss"] for h in t.state.log_history if "loss" in h]
calibration = FastDecisionModel.calibrate(model, processor, holdout) if holdout else {}
tuned = {name: _record_metrics(model, processor, its) for name, its in eval_items.items()}
result = {
"model": args.model,
"sources": sorted(pools) + (["typed_decisions"] if args.typed_decisions else []),
"train_rows": len(rows),
"train_records": len(train),
"holdout_records": len(holdout),
"skipped": report["skipped"],
"head": model.head.config,
"head_params_m": round(sum(p.numel() for p in model.head.parameters()) / 1e6, 1),
"head_warmup": warmup,
"steps": steps,
"epochs": args.epochs,
"loss_first_last": [round(losses[0], 4), round(losses[-1], 4)] if losses else None,
"s_per_step": round(out.metrics["train_runtime"] / max(1, steps), 3),
"peak_gb": round(torch.cuda.max_memory_allocated() / 1e9, 2)
if torch.cuda.is_available()
else None,
"temperatures": [round(x, 3) for x in model.decision_config["temperature"]],
"calibration_holdout": {k: calibration[k] for k in ("accuracy", "ece") if k in calibration},
"base": base,
"tuned": tuned,
"wall_min": round((time.time() - start) / 60, 1),
"gpu": torch.cuda.get_device_name() if torch.cuda.is_available() else "cpu",
}
if args.out:
model.save_pretrained_merged(args.out)
result["saved"] = args.out
line = json.dumps(result)
print("DECISION_FROM_LM_RESULT " + line)
if args.json:
with open(args.json, "w", encoding = "utf-8") as f:
f.write(line + "\n")
if __name__ == "__main__":
main()