| """Stage 2: continue training the saved stage-1 decision model (adapters + head) for broader reasoning. |
| |
| * typed-decisions train replay (half of it again) keeps the decision skills |
| * new commonsense / multiple-choice sources (none of them is an evaluation set): ARC-Easy, |
| OpenBookQA, COPA, WSC, Social IQa, PIQA, SWAG where loadable, plus fresh rows of Unsloth's public |
| mixture and of n4ze3m/typed-decisions-synth |
| * temperatures refit on a mixed holdout (typed-decisions holdout + unseen general items), so |
| confidences also hold outside typed-decisions |
| |
| Evaluation sets (typed-decisions test, ANLI, WinoGrande, HellaSwag, emotion) are never trained on. |
| No TypeSafe / Jev output is used. |
| """ |
|
|
| import argparse |
| import json |
| import os |
| import random |
| import shutil |
| import sys |
| import time |
|
|
| T0 = time.time() |
|
|
|
|
| def log(*parts): |
| print(f"[{time.strftime('%H:%M:%S')} +{(time.time() - T0) / 60:.1f}m]", *parts, flush = True) |
|
|
|
|
| parser = argparse.ArgumentParser() |
| parser.add_argument("--folder", default = "/content/out/final") |
| parser.add_argument("--base", default = "/content/models/Qwen3.8-27B") |
| parser.add_argument("--base-name", default = "Qwen/Qwen3.8-27B") |
| parser.add_argument("--max-seq-length", type = int, default = 2048) |
| parser.add_argument("--lr", type = float, default = 5e-5) |
| parser.add_argument("--head-lr", type = float, default = 1e-4) |
| parser.add_argument("--batch-size", type = int, default = 4) |
| parser.add_argument("--grad-accum", type = int, default = 4) |
| parser.add_argument("--train-minutes", type = float, default = 16.0) |
| parser.add_argument("--ms-per-token", type = float, default = 0.60) |
| parser.add_argument("--td-replay", type = int, default = 560) |
| parser.add_argument("--seed", type = int, default = 4242) |
| parser.add_argument("--out", default = "/content/out") |
| parser.add_argument("--name", default = "final2") |
| parser.add_argument("--skip-pre-eval", action = "store_true") |
| args = parser.parse_args() |
|
|
| |
| local = args.folder.rstrip("/") + "_local" |
| shutil.rmtree(local, ignore_errors = True) |
| shutil.copytree(args.folder, local) |
| for name, key in (("unsloth_decision_config.json", "base_model"), ("adapter_config.json", "base_model_name_or_path")): |
| path = os.path.join(local, name) |
| config = json.load(open(path, encoding = "utf-8")) |
| config[key] = args.base |
| json.dump(config, open(path, "w", encoding = "utf-8"), indent = 2) |
|
|
| from unsloth import DecisionTrainer, FastDecisionModel |
| from unsloth.models import decision_datasets as dd |
|
|
| import torch |
| from datasets import load_dataset |
| from transformers import TrainerCallback, TrainingArguments |
|
|
| sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) |
| from td_metrics import parsed, score_cases |
|
|
|
|
| class Progress(TrainerCallback): |
| def __init__(self, name, limit_s = None): |
| self.name, self.limit_s, self.start, self.steps = name, limit_s, None, [] |
|
|
| def on_train_begin(self, args, state, control, **kwargs): |
| self.start = time.time() |
| log(f"{self.name}: train begin, max_steps={state.max_steps}") |
|
|
| def on_step_end(self, args, state, control, **kwargs): |
| self.steps.append(time.time()) |
| if self.limit_s and time.time() - self.start > self.limit_s: |
| log(f"{self.name}: time limit reached at step {state.global_step}") |
| control.should_training_stop = True |
|
|
| def on_log(self, args, state, control, logs = None, **kwargs): |
| if logs: |
| recent = self.steps[-6:] |
| per_step = (recent[-1] - recent[0]) / (len(recent) - 1) if len(recent) > 1 else None |
| log(f"{self.name}: step {state.global_step}/{state.max_steps}", json.dumps(logs), |
| f"s/step={per_step:.2f}" if per_step else "") |
|
|
|
|
| def mcq(name, rows, question_key = "question"): |
| return list(dd.convert_mcq(name, question_key = question_key)(rows)) |
|
|
|
|
| def copa(rows): |
| out = [] |
| for row in rows: |
| if int(row["label"]) < 0: |
| continue |
| ask = "What was the cause?" if row["question"] == "cause" else "What happened as a result?" |
| out.append(dd._row( |
| "copa", {"premise": row["premise"], "question": ask}, "answer", |
| {"type": "choice", "instructions": "Which alternative is more plausible?", |
| "criteria": {"choice_1": row["choice1"], "choice_2": row["choice2"]}}, |
| "choice_1" if int(row["label"]) == 0 else "choice_2")) |
| return out |
|
|
|
|
| def wsc(rows): |
| out = [] |
| for row in rows: |
| if int(row["label"]) < 0: |
| continue |
| out.append(dd._row( |
| "wsc", {"sentence": row["text"]}, "refers", |
| {"type": "noul", "instructions": f"In the sentence, does \"{row['span2_text']}\" refer to \"{row['span1_text']}\"?"}, |
| bool(int(row["label"])))) |
| return out |
|
|
|
|
| def siqa(rows): |
| out = [] |
| for row in rows: |
| label = {"1": "A", "2": "B", "3": "C"}.get(str(row["label"]).strip()) |
| if not label: |
| continue |
| out.append(dd._row( |
| "social_i_qa", {"context": row["context"], "question": row["question"]}, "answer", |
| {"type": "choice", "instructions": "Which answer is most plausible?", |
| "criteria": {"A": row["answerA"], "B": row["answerB"], "C": row["answerC"]}}, label)) |
| return out |
|
|
|
|
| def piqa(rows): |
| out = [] |
| for row in rows: |
| if int(row["label"]) < 0: |
| continue |
| out.append(dd._row( |
| "piqa", {"goal": row["goal"]}, "solution", |
| {"type": "choice", "instructions": "Which solution achieves the goal?", |
| "criteria": {"solution_1": row["sol1"], "solution_2": row["sol2"]}}, |
| "solution_1" if int(row["label"]) == 0 else "solution_2")) |
| return out |
|
|
|
|
| def swag(rows): |
| out = [] |
| for row in rows: |
| if int(row["label"]) < 0: |
| continue |
| keys = ["A", "B", "C", "D"] |
| endings = [row[f"ending{i}"] for i in range(4)] |
| out.append(dd._row( |
| "swag", {"context": f"{row['sent1']} {row['sent2']}"}, "ending", |
| {"type": "choice", "instructions": "Which ending most plausibly continues the context?", |
| "criteria": dict(zip(keys, endings))}, keys[int(row["label"])])) |
| return out |
|
|
|
|
| SOURCES = [ |
| ("arc_easy", lambda: mcq("arc_easy", load_dataset("allenai/ai2_arc", "ARC-Easy", split = "train"))), |
| ("openbookqa", lambda: mcq("openbookqa", load_dataset("allenai/openbookqa", "main", split = "train"), "question_stem")), |
| ("copa", lambda: copa(load_dataset("aps/super_glue", "copa", split = "train"))), |
| ("wsc", lambda: wsc(load_dataset("aps/super_glue", "wsc.fixed", split = "train"))), |
| ("social_i_qa", lambda: siqa(load_dataset("allenai/social_i_qa", split = "train"))), |
| ("piqa", lambda: piqa(load_dataset("ybisk/piqa", split = "train"))), |
| ("swag", lambda: swag(load_dataset("allenai/swag", "regular", split = "train"))), |
| ] |
|
|
|
|
| def predict_cases(model, tokenizer, rows, name): |
| FastDecisionModel.for_inference(model) |
| cases = [] |
| for row in rows: |
| answers = FastDecisionModel.predict(model, tokenizer, parsed(row["state"]), parsed(row["questions"])) |
| cases.append({"id": row["id"], "workflow": row["workflow"], "questions": row["questions"], "gold": row["gold"], |
| "pred": {q: {k: v for k, v in a.items() if k != "legend"} for q, a in answers.items()}}) |
| metrics = score_cases(cases) |
| with open(f"{args.out}/{name}_preds.jsonl", "w", encoding = "utf-8") as f: |
| for case in cases: |
| f.write(json.dumps(case) + "\n") |
| with open(f"{args.out}/{name}_metrics.json", "w", encoding = "utf-8") as f: |
| json.dump(metrics, f, indent = 1) |
| log(f"{name.upper()}_METRICS", json.dumps(metrics)) |
| return metrics |
|
|
|
|
| def jsonl(path): |
| return [json.loads(line) for line in open(path, encoding = "utf-8") if line.strip()] |
|
|
|
|
| def main(): |
| log("args", json.dumps(vars(args))) |
| model, tokenizer = FastDecisionModel.from_pretrained(local, max_seq_length = args.max_seq_length) |
| model.decision_config["base_model"] = args.base_name |
| trainable = sum(p.numel() for p in model.parameters() if p.requires_grad) |
| if trainable < 1e6: |
| for name, p in model.named_parameters(): |
| if "lora_" in name or name.startswith("head."): |
| p.requires_grad_(True) |
| trainable = sum(p.numel() for p in model.parameters() if p.requires_grad) |
| log("loaded stage 1", json.dumps(model.decision_config), f"trainable={trainable / 1e6:.1f}M") |
|
|
| td_test = list(load_dataset("LocalLLaMA/typed-decisions", "all", split = "test")) |
| heldout = jsonl(f"{args.out}/heldout_rows.jsonl") |
| heldout2 = jsonl(f"{args.out}/heldout2_rows.jsonl") |
| if not args.skip_pre_eval: |
| predict_cases(model, tokenizer, td_test, "stage1_saved_td_test") |
|
|
| eval_texts = [parsed(r["state"]) for r in td_test + heldout + heldout2] |
| cleaner = dd.Decontaminator(eval_texts) |
| test_states = {json.dumps(parsed(r["state"]), sort_keys = True) for r in td_test} |
| rng = random.Random(3407 + 2) |
| td_rows = [dd.augment_row(r, rng) for r in dd.load_source("typed_decisions", "train") |
| if json.dumps(r["state"], sort_keys = True) not in test_states] |
| td_items, _ = FastDecisionModel.build_dataset(td_rows, tokenizer, model) |
| td_train, td_holdout = FastDecisionModel.split_holdout(td_items, 3407) |
|
|
| general = [] |
| for name, loader in SOURCES: |
| try: |
| rows = [r for r in loader() if not cleaner.contaminated(r)] |
| random.Random(args.seed).shuffle(rows) |
| general.append((name, rows[:700])) |
| log(f"source {name}: {len(rows)} rows") |
| except Exception as exc: |
| log(f"skipping {name}: {type(exc).__name__}: {str(exc)[:160]}") |
| rng = random.Random(args.seed) |
| new_rows = [dd.augment_row(r, rng) for _, rows in general for r in rows] |
| public = dd.build_decision_mixture(list(("banking77", "clinc150", "mnli", "snli", "wanli", "boolq", "ag_news", |
| "sst5", "mmlu", "commonsense_qa", "arc", "prompt_injections")), |
| n_rows = 1200, seed = args.seed, decontaminate_against = eval_texts) |
| synth = [] |
| for row in load_dataset("n4ze3m/typed-decisions-synth", split = "train"): |
| questions, teacher = parsed(row["questions"]), parsed(row["teacher"]) |
| if not isinstance(questions, dict) or not isinstance(teacher, dict): |
| continue |
| gold = {} |
| for qid, q in questions.items(): |
| a = teacher.get(qid) |
| if isinstance(a, dict) and q.get("type") == "noul" and a.get("noul") is not None: |
| gold[qid] = {"probabilities": {"false": 1 - float(a["noul"]), "true": float(a["noul"])}} |
| elif isinstance(a, dict) and a.get("probabilities"): |
| gold[qid] = {"probabilities": a["probabilities"]} |
| if gold: |
| state = parsed(row["state"]) if row.get("state_is_json") else row["state"] |
| synth.append({"state": state, "questions": {q: questions[q] for q in gold}, "gold": gold}) |
| random.Random(args.seed + 1).shuffle(synth) |
| synth = [r for r in synth[:900] if not cleaner.contaminated(r)] |
| general_items, _ = FastDecisionModel.build_dataset(new_rows + public + synth, tokenizer, model) |
| random.Random(args.seed + 2).shuffle(general_items) |
| calib_general, general_items = general_items[:320], general_items[320:] |
| log(f"data: td_train={len(td_train)} td_holdout={len(td_holdout)} general={len(general_items)} calib_general={len(calib_general)}") |
|
|
| replay = random.Random(args.seed + 3).sample(td_train, min(args.td_replay, len(td_train))) |
| budget = args.train_minutes * 60 / (args.ms_per_token / 1000) |
| selected, used = list(replay), sum(len(x["input_ids"]) for x in replay) |
| for item in general_items: |
| if used >= budget: |
| break |
| selected.append(item) |
| used += len(item["input_ids"]) |
| random.Random(args.seed + 4).shuffle(selected) |
| per_step = args.batch_size * args.grad_accum |
| steps = max(1, len(selected) // per_step) |
| log(f"selected {len(selected)} items ({used / 1e6:.2f}M tokens) -> {steps} steps") |
|
|
| FastDecisionModel.for_training(model) |
| trainer = DecisionTrainer( |
| model = model, tokenizer = tokenizer, train_dataset = selected, head_learning_rate = args.head_lr, |
| args = TrainingArguments( |
| output_dir = f"{args.out}/_trainer2", 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)), weight_decay = 0.01, max_steps = steps, |
| bf16 = torch.cuda.is_bf16_supported(), logging_steps = 5, report_to = "none", save_strategy = "no", |
| seed = args.seed, dataloader_num_workers = 0), |
| callbacks = [Progress("stage2", limit_s = args.train_minutes * 60 * 1.15)], |
| ) |
| out = trainer.train() |
| log("stage2 train done", json.dumps(out.metrics)) |
|
|
| calibration = FastDecisionModel.calibrate(model, tokenizer, td_holdout + calib_general) |
| calibration = {k: v for k, v in calibration.items() if not isinstance(v, (list, dict))} |
| log("calibration (mixed)", json.dumps(calibration), json.dumps(model.decision_config)) |
|
|
| metrics = { |
| "typed_decisions_test": predict_cases(model, tokenizer, td_test, f"{args.name}_td_test"), |
| "heldout": predict_cases(model, tokenizer, heldout, f"{args.name}_heldout"), |
| "heldout2": predict_cases(model, tokenizer, heldout2, f"{args.name}_heldout2"), |
| } |
| stage1 = json.load(open(f"{args.folder}/train_report.json", encoding = "utf-8")) |
| report = { |
| **stage1, |
| "stage2": { |
| "train_items": len(selected), "train_tokens": used, "steps": trainer.state.global_step, |
| "td_replay": len(replay), "sources": {name: len(rows) for name, rows in general}, |
| "train_metrics": out.metrics, "calibration_mixed_holdout": calibration, "args": vars(args), |
| }, |
| "train_items": stage1["train_items"] + len(selected), |
| "train_tokens": stage1["train_tokens"] + used, |
| "steps": stage1["steps"] + trainer.state.global_step, |
| "temperatures": model.decision_config.get("temperature"), |
| "eval": {"typed_decisions_test": metrics["typed_decisions_test"], "heldout": metrics["heldout"]}, |
| } |
| folder = f"{args.out}/{args.name}" |
| model.decision_config["base_model"] = args.base_name |
| model.save_pretrained(folder, tokenizer) |
| readme = os.path.join(folder, "README.md") |
| if os.path.exists(readme): |
| os.remove(readme) |
| with open(f"{folder}/train_report.json", "w", encoding = "utf-8") as f: |
| json.dump(report, f, indent = 1) |
| log(f"saved {folder}", sorted(os.listdir(folder))) |
| log("STAGE2_DONE", json.dumps({k: v["accuracy"] for k, v in metrics.items()})) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
