code/train_stage2.py

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

# Loadable copy of the stage-1 folder whose configs name the local base.
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  # noqa: E402
from unsloth.models import decision_datasets as dd  # noqa: E402

import torch  # noqa: E402
from datasets import load_dataset  # noqa: E402
from transformers import TrainerCallback, TrainingArguments  # noqa: E402

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from td_metrics import parsed, score_cases  # noqa: E402


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:  # adapters loaded frozen: make LoRA and the head trainable again
        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:  # the saved stage-1 artifact must reproduce its in-process numbers
        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)  # the stage-1 augmentation and split, so the same 80 rows stay held out
    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)  # stage 1 drew its 3,000 with seed 3407; these are mostly new
    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)  # PEFT's card names the local base path; finalize.py writes the real card
    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()