Buckets:

Mercity
/

policy-classifier

Files

xet

Mercity/policy-classifier / 05-11-progress /training /engine.py

Rishik001

3 days ago

download

raw

22 kB

"""Shared training machinery for both classifier scripts (train_modernbert.py, train_qwen.py).
Each model script only decides how to load its backbone, which modules get LoRA, which module is the
classification head, and its data collator. Everything else lives here so both models are trained,
checkpointed and measured identically:
  * dataset validation (labels, empty text, case/input leakage between splits);
  * LoRA/DoRA attachment with the classification head trained in full;
  * tokenization in the shared input format (training.common.tokenize_pairs);
  * adapter-only checkpoints at 1/4, 2/4, 3/4 and the end of training, each followed by evaluation;
  * scikit-learn metrics, a confusion matrix and per-row validation predictions.
ML imports are kept inside functions so `--preflight` runs without torch installed.
"""

import json
import time
from collections import Counter
from pathlib import Path

import yaml

from training.common import (
    INSTRUCTION,
    LABEL2ID,
    LABELS,
    classification_metrics,
    confusion_matrix_counts,
    read_rows,
    tokenize_pairs,
    write_rows,
)

REQUIRED_KEYS = (
    "model_name", "data_dir", "output_dir", "run_name", "wandb_project", "seed", "max_length",
    "epochs", "learning_rate", "lr_scheduler_type", "warmup_ratio", "weight_decay", "max_grad_norm",
    "batch_size", "gradient_accumulation_steps", "eval_batch_size", "gradient_checkpointing",
    "attention_implementation", "lora_rank", "lora_alpha", "lora_dropout", "use_dora",
    "checkpoints_per_run", "logging_steps", "dataloader_num_workers",
)


def load_config(path, epochs=None):
    config = yaml.safe_load(Path(path).read_text(encoding="utf-8"))
    missing = [key for key in REQUIRED_KEYS if key not in config]
    if missing:
        raise ValueError(f"Config {path} is missing: {', '.join(missing)}")
    if epochs is not None:
        config["epochs"] = epochs
    return config


def validate_inputs(config):
    """Check the config and the dataset before any model is downloaded."""
    if not 0 < config["max_length"] <= 8192:
        raise ValueError("The configured classifier training cap must be between 1 and 8192 tokens")
    if config["epochs"] <= 0 or config["batch_size"] <= 0 or config["gradient_accumulation_steps"] <= 0:
        raise ValueError("Epochs and batch sizes must be positive")
    if config.get("lora_rank", 1) <= 0:
        raise ValueError("LoRA rank must be positive")
    data_dir = Path(config["data_dir"])
    splits = {name: read_rows(data_dir / f"{name}.jsonl") for name in ("train", "validation")}
    if (data_dir / "test_ood.jsonl").exists():
        splits["test_ood"] = read_rows(data_dir / "test_ood.jsonl")
    cases, inputs = {}, {}
    for split, rows in splits.items():
        if split != "test_ood" and not rows:
            raise ValueError(f"{split} split is empty")
        for row in rows:
            if row["label"] not in LABEL2ID or row["labels"] != LABEL2ID[row["label"]]:
                raise ValueError("Dataset label ids disagree with the training label map")
            if not row["message"].strip() or not row["policy"].strip():
                raise ValueError("Empty message or policy")
            if split != "test_ood" and row["held_out"]:
                raise ValueError("Held-out domain present in train/validation")
            if row["case_id"] in cases and cases[row["case_id"]] != split:
                raise ValueError("A case appears in multiple splits")
            key = (row["message"], row["policy"])
            if key in inputs and inputs[key] != split:
                raise ValueError("Identical input appears in multiple splits")
            cases[row["case_id"]], inputs[key] = split, split
    if set(row["label"] for row in splits["train"]) != set(LABELS):
        raise ValueError("Training requires all three classes")
    return splits


def describe_splits(splits):
    return {
        name: {"pairs": len(rows), "labels": dict(Counter(row["label"] for row in rows))}
        for name, rows in splits.items()
    }


def classifier_config(base_config, tokenizer):
    """Standard HF three-class single-label setup (cross-entropy loss on the classification head)."""
    base_config.num_labels = len(LABELS)
    base_config.label2id, base_config.id2label = LABEL2ID, dict(enumerate(LABELS))
    base_config.problem_type = "single_label_classification"
    base_config.pad_token_id = tokenizer.pad_token_id
    return base_config


def attach_lora(model, config, target_modules, head_modules):
    """LoRA (or DoRA when use_dora) on the backbone projections; the new head is trained in full and
    saved inside the adapter checkpoint through modules_to_save."""
    from peft import LoraConfig, TaskType, get_peft_model

    adapter = LoraConfig(
        task_type=TaskType.SEQ_CLS,
        r=config["lora_rank"],
        lora_alpha=config["lora_alpha"],
        lora_dropout=config["lora_dropout"],
        use_dora=config["use_dora"],
        bias="none",
        target_modules=target_modules,
        modules_to_save=head_modules,
    )
    model = get_peft_model(model, adapter)
    upcast_trainable_parameters(model)
    model.print_trainable_parameters()
    return model


def upcast_trainable_parameters(model):
    """Keep every trainable weight in float32; the frozen backbone stays in BF16.
    The backbone is loaded in BF16. PEFT creates the LoRA weights in float32 but leaves the head copies
    (modules_to_save) in BF16, where most of Adam's small updates round to zero: on the 30k runs 72% of
    ModernBERT's dense head weights never moved. BF16 autocast (Trainer bf16=True) still runs the head's
    matmuls in BF16, so this costs almost nothing in speed; only the stored weights and updates are float32.
    Returns the names of the parameters that were converted.
    """
    import torch

    converted = []
    for name, parameter in model.named_parameters():
        if parameter.requires_grad and parameter.dtype != torch.float32:
            parameter.data = parameter.data.float()
            converted.append(name)
    return converted


def seed_everything(seed):
    """Seed Python, NumPy and torch. Call before building the model so the new head and the LoRA
    weights start from the same values on every run with the same seed."""
    from transformers import set_seed

    set_seed(seed)


def input_recipe(config, tokenizer, base_config):
    """Everything needed to rebuild the exact classifier input at prediction time."""
    return {
        "model_name": config["model_name"],
        "model_revision": getattr(base_config, "_commit_hash", None),
        "model_type": base_config.model_type,
        "max_length": config["max_length"],
        "instruction": INSTRUCTION,
        "label2id": LABEL2ID,
        "separator": tokenizer.sep_token or "\n\n",
        "pad_token_id": tokenizer.pad_token_id,
        "attention_implementation": config["attention_implementation"],
        "format": "tokenizer(text=instruction+message+optional_blank_line, text_pair=policy)",
    }


def save_adapter(model, tokenizer, recipe, path):
    """Adapter weights + trained classification head only; no optimizer or full model weights."""
    path = Path(path)
    model.save_pretrained(path)
    tokenizer.save_pretrained(path)
    (path / "input_format.json").write_text(json.dumps(recipe, indent=2), encoding="utf-8")


def quarter_checkpoint_callback(tokenizer, recipe, output, count):
    """Save the adapter and run evaluation at k/count of the total optimizer steps, k = 1..count."""
    from transformers import TrainerCallback

    class EvenCheckpoints(TrainerCallback):
        def on_train_begin(self, args, state, control, **kwargs):
            # state.max_steps is final here; the last checkpoint always lands on the last step.
            self.steps = {max(1, round(state.max_steps * k / count)): k for k in range(1, count + 1)}

        def on_step_end(self, args, state, control, model=None, **kwargs):
            part = self.steps.get(state.global_step)
            if part is None:
                return control
            control.should_evaluate = True  # Trainer evaluates right after this callback
            if state.is_world_process_zero:
                path = Path(output) / f"checkpoint-{part}of{count}-step{state.global_step}"
                save_adapter(model, tokenizer, recipe, path)
                print(f"Saved adapter checkpoint {path}", flush=True)
            return control

    return EvenCheckpoints()


def prediction_logits(prediction):
    """Trainer may return logits alone or a tuple containing logits and auxiliary outputs."""
    return prediction.predictions[0] if isinstance(prediction.predictions, tuple) else prediction.predictions


def compute_metrics(prediction):
    return classification_metrics(
        prediction.label_ids.tolist(), prediction_logits(prediction).argmax(axis=-1).tolist()
    )


def softmax_rows(logits):
    import numpy as np

    shifted = logits - logits.max(axis=-1, keepdims=True)
    exp = np.exp(shifted)
    return exp / exp.sum(axis=-1, keepdims=True)


def prediction_rows(rows, logits):
    """One record per input with the gold label, predicted label and class probabilities."""
    probabilities = softmax_rows(logits.astype("float64"))
    keep = ("pair_id", "case_id", "domain", "policy_id", "conv_id", "synthetic", "selection_tier",
            "message", "policy", "label")
    return [
        {
            **{key: row.get(key) for key in keep},
            "predicted": LABELS[int(probs.argmax())],
            "probabilities": {name: float(probs[index]) for index, name in enumerate(LABELS)},
        }
        for row, probs in zip(rows, probabilities, strict=True)
    ]


def evaluation_report(records):
    """Metrics, confusion matrix and the human-readable sklearn report from prediction records.
    `by_source` repeats metrics and confusion matrix for natural pairs (original replies) and edited
    pairs (replies changed by the edit stage) separately. Most training violations are edits, so the
    overall numbers can look good while natural violations are missed; the natural rows are the ones
    that resemble production.
    """
    from sklearn.metrics import classification_report

    truth = [LABEL2ID[record["label"]] for record in records]
    predicted = [LABEL2ID[record["predicted"]] for record in records]
    by_source = {}
    for source, is_edited in (("natural", False), ("edited", True)):
        subset = [record for record in records if bool(record.get("synthetic")) == is_edited]
        if subset:
            subset_truth = [LABEL2ID[record["label"]] for record in subset]
            subset_predicted = [LABEL2ID[record["predicted"]] for record in subset]
            by_source[source] = {
                "pairs": len(subset),
                "metrics": classification_metrics(subset_truth, subset_predicted),
                "confusion_matrix": confusion_matrix_counts(subset_truth, subset_predicted),
            }
    return {
        "metrics": classification_metrics(truth, predicted),
        "confusion_matrix": confusion_matrix_counts(truth, predicted),
        "by_source": by_source,
        "classification_report": classification_report(
            truth, predicted, labels=list(range(len(LABELS))), target_names=list(LABELS),
            digits=4, zero_division=0,
        ),
    }


def source_metrics_for_logging(split, report):
    """Flat W&B keys such as validation_natural/violated_recall from evaluation_report's by_source."""
    return {
        f"{split}_{source}/{name}": value
        for source, part in report["by_source"].items()
        for name, value in part["metrics"].items()
    }


def require_online_wandb(config):
    """Fail before training if W&B would log offline. The 30k runs logged offline (set by the host's
    environment), so they had no run URLs until synced by hand. Set allow_offline_wandb: true in the
    config to permit it on a host without network access."""
    import os

    mode = os.environ.get("WANDB_MODE", "online").lower()
    if mode in {"offline", "dryrun", "disabled"} and not config.get("allow_offline_wandb", False):
        raise RuntimeError(
            f"WANDB_MODE={mode}: runs would not appear in W&B. Unset WANDB_MODE (and log in with "
            "WANDB_API_KEY), or set allow_offline_wandb: true in the config."
        )


def run_training(config, model, tokenizer, base_config, collator, resume_from_checkpoint=None):
    """Train with HF Trainer and W&B logging; adapters saved at every checkpoint fraction."""
    import torch
    import wandb
    from transformers import Trainer, TrainingArguments, set_seed

    from datasets import Dataset

    if not torch.cuda.is_available() or not torch.cuda.is_bf16_supported():
        raise RuntimeError("Training requires a CUDA GPU with BF16 support (configured for A100)")
    require_online_wandb(config)
    splits = validate_inputs(config)
    set_seed(config["seed"])

    def tokenize(batch):
        encoded = tokenize_pairs(tokenizer, batch["message"], batch["policy"], config["max_length"])
        encoded["labels"] = batch["labels"]
        return encoded

    datasets = {
        name: Dataset.from_list(rows).map(tokenize, batched=True, remove_columns=list(rows[0]))
        for name, rows in splits.items()
        if rows
    }
    output = Path(config["output_dir"])
    output.mkdir(parents=True, exist_ok=True)
    (output / "training_config.json").write_text(json.dumps(config, indent=2), encoding="utf-8")
    recipe = input_recipe(config, tokenizer, base_config)
    # transformers 5 folded warmup_ratio into warmup_steps (a float < 1 is a ratio) and dropped
    # group_by_length; 4.57 keeps the original arguments, so behaviour there is unchanged.
    fields = TrainingArguments.__dataclass_fields__
    version_args = {}
    if "warmup_ratio" in fields:
        version_args["warmup_ratio"] = config["warmup_ratio"]
    else:
        version_args["warmup_steps"] = float(config["warmup_ratio"])
    if "group_by_length" in fields:
        version_args["group_by_length"] = config.get("group_by_length", False)
    elif config.get("group_by_length", False):
        raise ValueError("group_by_length is not supported by this transformers version")
    training_args = TrainingArguments(
        **version_args,
        output_dir=str(output),
        run_name=config["run_name"],
        num_train_epochs=config["epochs"],
        learning_rate=config["learning_rate"],
        lr_scheduler_type=config["lr_scheduler_type"],
        weight_decay=config["weight_decay"],
        max_grad_norm=config["max_grad_norm"],
        per_device_train_batch_size=config["batch_size"],
        per_device_eval_batch_size=config["eval_batch_size"],
        gradient_accumulation_steps=config["gradient_accumulation_steps"],
        bf16=True,
        tf32=True,
        gradient_checkpointing=config["gradient_checkpointing"],
        gradient_checkpointing_kwargs={"use_reentrant": False},
        # Evaluation and saving are driven by the checkpoint callback, never by Trainer itself.
        eval_strategy="no",
        save_strategy="no",
        logging_steps=config["logging_steps"],
        report_to=["wandb"],
        seed=config["seed"],
        data_seed=config["seed"],
        optim="adamw_torch",
        dataloader_num_workers=config["dataloader_num_workers"],
        eval_accumulation_steps=8,
    )
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=datasets["train"],
        eval_dataset=datasets["validation"],
        processing_class=tokenizer,
        data_collator=collator,
        compute_metrics=compute_metrics,
        callbacks=[quarter_checkpoint_callback(tokenizer, recipe, output, config["checkpoints_per_run"])],
    )
    # Open the W&B run only after setup succeeds; the finally block closes it on errors too.
    run = wandb.init(project=config["wandb_project"], name=config["run_name"], config=config)
    (output / "wandb_run.json").write_text(json.dumps({"url": run.url, "id": run.id}), encoding="utf-8")
    try:
        torch.cuda.reset_peak_memory_stats()
        started = time.perf_counter()
        trained = trainer.train(resume_from_checkpoint=resume_from_checkpoint)
        torch.cuda.synchronize()
        hardware = {
            "gpu": torch.cuda.get_device_name(),
            "peak_allocated_gib": torch.cuda.max_memory_allocated() / 2**30,
            "peak_reserved_gib": torch.cuda.max_memory_reserved() / 2**30,
            "training_wall_seconds": time.perf_counter() - started,
        }
        (output / "hardware_metrics.json").write_text(json.dumps(hardware, indent=2), encoding="utf-8")
        wandb.log(hardware)
        trainer.save_metrics("train", trained.metrics)
        trainer.save_state()
        save_adapter(model, tokenizer, recipe, output / "final")  # same weights as the last checkpoint
        for split in ("validation", "test_ood"):
            if split not in datasets:
                continue
            prediction = trainer.predict(datasets[split], metric_key_prefix=split)
            records = prediction_rows(splits[split], prediction_logits(prediction))
            report = evaluation_report(records)
            write_rows(output / f"{split}_predictions.jsonl", records)
            (output / f"{split}_report.json").write_text(json.dumps(report, indent=2), encoding="utf-8")
            print(report["classification_report"], flush=True)
            for source, part in report["by_source"].items():
                metrics = part["metrics"]
                print(f"{split} {source} ({part['pairs']} pairs): macro F1 {metrics['macro_f1']:.4f}, "
                      f"violated precision {metrics['violated_precision']:.4f}, "
                      f"violated recall {metrics['violated_recall']:.4f}", flush=True)
            trainer.log_metrics(split, prediction.metrics)
            wandb.log(prediction.metrics)
            wandb.log(source_metrics_for_logging(split, report))
            wandb.log(
                {
                    f"{split}_confusion_matrix": wandb.plot.confusion_matrix(
                        y_true=[LABEL2ID[r["label"]] for r in records],
                        preds=[LABEL2ID[r["predicted"]] for r in records],
                        class_names=list(LABELS),
                    ),
                    f"{split}_errors": wandb.Table(
                        columns=["message", "policy", "actual", "predicted", "p_violated"],
                        data=[
                            [r["message"], r["policy"], r["label"], r["predicted"],
                             r["probabilities"]["violated"]]
                            for r in records if r["label"] != r["predicted"]
                        ][:500],
                    ),
                }
            )
    finally:
        wandb.finish()


def load_trained_classifier(checkpoint, device):
    """Rebuild base model + saved adapter/head from a checkpoint directory's input_format.json."""
    import torch
    from peft import PeftModel
    from transformers import AutoModelForSequenceClassification, AutoTokenizer

    checkpoint = Path(checkpoint)
    recipe = json.loads((checkpoint / "input_format.json").read_text(encoding="utf-8"))
    tokenizer = AutoTokenizer.from_pretrained(checkpoint)
    dtype = torch.bfloat16 if device == "cuda" and torch.cuda.is_bf16_supported() else torch.float32
    extra = {"reference_compile": False} if recipe["model_type"] == "modernbert" else {"use_cache": False}
    if recipe["model_type"] == "qwen3_5_text":
        # Qwen3.5 checkpoints are vision-language; the classifier uses only the text decoder.
        from training.train_qwen35 import load_text_classifier

        base = load_text_classifier(
            recipe["model_name"], recipe.get("model_revision"), recipe["label2id"], tokenizer.pad_token_id,
            dtype, recipe.get("attention_implementation", "sdpa"),
        )
        model = PeftModel.from_pretrained(base, checkpoint).to(device).eval()
        return model, tokenizer, recipe
    base = AutoModelForSequenceClassification.from_pretrained(
        recipe["model_name"],
        revision=recipe.get("model_revision"),
        num_labels=len(recipe["label2id"]),
        label2id=recipe["label2id"],
        id2label={index: label for label, index in recipe["label2id"].items()},
        pad_token_id=tokenizer.pad_token_id,
        torch_dtype=dtype,
        attn_implementation=recipe.get("attention_implementation", "sdpa"),
        **extra,
    )
    model = PeftModel.from_pretrained(base, checkpoint).to(device).eval()
    return model, tokenizer, recipe


def predict_logits(model, tokenizer, rows, recipe, batch_size, device):
    """Batched inference in the saved input format; returns a numpy array of logits."""
    import numpy as np
    import torch

    # Longest-first batches waste the least padding; results are put back in input order.
    order = sorted(range(len(rows)), key=lambda index: -len(rows[index]["message"]) - len(rows[index]["policy"]))
    logits = [None] * len(rows)
    with torch.inference_mode():
        for start in range(0, len(order), batch_size):
            chunk = order[start : start + batch_size]
            batch = tokenize_pairs(
                tokenizer,
                [rows[index]["message"] for index in chunk],
                [rows[index]["policy"] for index in chunk],
                recipe["max_length"],
                recipe["instruction"],
                padding=True,
                return_tensors="pt",
            ).to(device)
            output = model(**batch).logits.float().cpu().numpy()
            for index, value in zip(chunk, output, strict=True):
                logits[index] = value
            print(f"predicted {min(start + batch_size, len(order))}/{len(order)}", flush=True)
    return np.stack(logits)

Xet Storage Details

Size:: 22 kB
Xet hash:: 12c86430f91fb58926edf11394313ab78841c27a700f1aba76b1c485974deb4a

Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.

	"""Shared training machinery for both classifier scripts (train_modernbert.py, train_qwen.py).

	Each model script only decides how to load its backbone, which modules get LoRA, which module is the
	classification head, and its data collator. Everything else lives here so both models are trained,
	checkpointed and measured identically:
	* dataset validation (labels, empty text, case/input leakage between splits);
	* LoRA/DoRA attachment with the classification head trained in full;
	* tokenization in the shared input format (training.common.tokenize_pairs);
	* adapter-only checkpoints at 1/4, 2/4, 3/4 and the end of training, each followed by evaluation;
	* scikit-learn metrics, a confusion matrix and per-row validation predictions.

	ML imports are kept inside functions so `--preflight` runs without torch installed.
	"""

	import json
	import time
	from collections import Counter
	from pathlib import Path

	import yaml

	from training.common import (
	INSTRUCTION,
	LABEL2ID,
	LABELS,
	classification_metrics,
	confusion_matrix_counts,
	read_rows,
	tokenize_pairs,
	write_rows,
	)

	REQUIRED_KEYS = (
	"model_name", "data_dir", "output_dir", "run_name", "wandb_project", "seed", "max_length",
	"epochs", "learning_rate", "lr_scheduler_type", "warmup_ratio", "weight_decay", "max_grad_norm",
	"batch_size", "gradient_accumulation_steps", "eval_batch_size", "gradient_checkpointing",
	"attention_implementation", "lora_rank", "lora_alpha", "lora_dropout", "use_dora",
	"checkpoints_per_run", "logging_steps", "dataloader_num_workers",
	)


	def load_config(path, epochs=None):
	config = yaml.safe_load(Path(path).read_text(encoding="utf-8"))
	missing = [key for key in REQUIRED_KEYS if key not in config]
	if missing:
	raise ValueError(f"Config {path} is missing: {', '.join(missing)}")
	if epochs is not None:
	config["epochs"] = epochs
	return config


	def validate_inputs(config):
	"""Check the config and the dataset before any model is downloaded."""
	if not 0 < config["max_length"] <= 8192:
	raise ValueError("The configured classifier training cap must be between 1 and 8192 tokens")
	if config["epochs"] <= 0 or config["batch_size"] <= 0 or config["gradient_accumulation_steps"] <= 0:
	raise ValueError("Epochs and batch sizes must be positive")
	if config.get("lora_rank", 1) <= 0:
	raise ValueError("LoRA rank must be positive")
	data_dir = Path(config["data_dir"])
	splits = {name: read_rows(data_dir / f"{name}.jsonl") for name in ("train", "validation")}
	if (data_dir / "test_ood.jsonl").exists():
	splits["test_ood"] = read_rows(data_dir / "test_ood.jsonl")
	cases, inputs = {}, {}
	for split, rows in splits.items():
	if split != "test_ood" and not rows:
	raise ValueError(f"{split} split is empty")
	for row in rows:
	if row["label"] not in LABEL2ID or row["labels"] != LABEL2ID[row["label"]]:
	raise ValueError("Dataset label ids disagree with the training label map")
	if not row["message"].strip() or not row["policy"].strip():
	raise ValueError("Empty message or policy")
	if split != "test_ood" and row["held_out"]:
	raise ValueError("Held-out domain present in train/validation")
	if row["case_id"] in cases and cases[row["case_id"]] != split:
	raise ValueError("A case appears in multiple splits")
	key = (row["message"], row["policy"])
	if key in inputs and inputs[key] != split:
	raise ValueError("Identical input appears in multiple splits")
	cases[row["case_id"]], inputs[key] = split, split
	if set(row["label"] for row in splits["train"]) != set(LABELS):
	raise ValueError("Training requires all three classes")
	return splits


	def describe_splits(splits):
	return {
	name: {"pairs": len(rows), "labels": dict(Counter(row["label"] for row in rows))}
	for name, rows in splits.items()
	}


	def classifier_config(base_config, tokenizer):
	"""Standard HF three-class single-label setup (cross-entropy loss on the classification head)."""
	base_config.num_labels = len(LABELS)
	base_config.label2id, base_config.id2label = LABEL2ID, dict(enumerate(LABELS))
	base_config.problem_type = "single_label_classification"
	base_config.pad_token_id = tokenizer.pad_token_id
	return base_config


	def attach_lora(model, config, target_modules, head_modules):
	"""LoRA (or DoRA when use_dora) on the backbone projections; the new head is trained in full and
	saved inside the adapter checkpoint through modules_to_save."""
	from peft import LoraConfig, TaskType, get_peft_model

	adapter = LoraConfig(
	task_type=TaskType.SEQ_CLS,
	r=config["lora_rank"],
	lora_alpha=config["lora_alpha"],
	lora_dropout=config["lora_dropout"],
	use_dora=config["use_dora"],
	bias="none",
	target_modules=target_modules,
	modules_to_save=head_modules,
	)
	model = get_peft_model(model, adapter)
	upcast_trainable_parameters(model)
	model.print_trainable_parameters()
	return model


	def upcast_trainable_parameters(model):
	"""Keep every trainable weight in float32; the frozen backbone stays in BF16.

	The backbone is loaded in BF16. PEFT creates the LoRA weights in float32 but leaves the head copies
	(modules_to_save) in BF16, where most of Adam's small updates round to zero: on the 30k runs 72% of
	ModernBERT's dense head weights never moved. BF16 autocast (Trainer bf16=True) still runs the head's
	matmuls in BF16, so this costs almost nothing in speed; only the stored weights and updates are float32.
	Returns the names of the parameters that were converted.
	"""
	import torch

	converted = []
	for name, parameter in model.named_parameters():
	if parameter.requires_grad and parameter.dtype != torch.float32:
	parameter.data = parameter.data.float()
	converted.append(name)
	return converted


	def seed_everything(seed):
	"""Seed Python, NumPy and torch. Call before building the model so the new head and the LoRA
	weights start from the same values on every run with the same seed."""
	from transformers import set_seed

	set_seed(seed)


	def input_recipe(config, tokenizer, base_config):
	"""Everything needed to rebuild the exact classifier input at prediction time."""
	return {
	"model_name": config["model_name"],
	"model_revision": getattr(base_config, "_commit_hash", None),
	"model_type": base_config.model_type,
	"max_length": config["max_length"],
	"instruction": INSTRUCTION,
	"label2id": LABEL2ID,
	"separator": tokenizer.sep_token or "\n\n",
	"pad_token_id": tokenizer.pad_token_id,
	"attention_implementation": config["attention_implementation"],
	"format": "tokenizer(text=instruction+message+optional_blank_line, text_pair=policy)",
	}


	def save_adapter(model, tokenizer, recipe, path):
	"""Adapter weights + trained classification head only; no optimizer or full model weights."""
	path = Path(path)
	model.save_pretrained(path)
	tokenizer.save_pretrained(path)
	(path / "input_format.json").write_text(json.dumps(recipe, indent=2), encoding="utf-8")


	def quarter_checkpoint_callback(tokenizer, recipe, output, count):
	"""Save the adapter and run evaluation at k/count of the total optimizer steps, k = 1..count."""
	from transformers import TrainerCallback

	class EvenCheckpoints(TrainerCallback):
	def on_train_begin(self, args, state, control, **kwargs):
	# state.max_steps is final here; the last checkpoint always lands on the last step.
	self.steps = {max(1, round(state.max_steps * k / count)): k for k in range(1, count + 1)}

	def on_step_end(self, args, state, control, model=None, **kwargs):
	part = self.steps.get(state.global_step)
	if part is None:
	return control
	control.should_evaluate = True # Trainer evaluates right after this callback
	if state.is_world_process_zero:
	path = Path(output) / f"checkpoint-{part}of{count}-step{state.global_step}"
	save_adapter(model, tokenizer, recipe, path)
	print(f"Saved adapter checkpoint {path}", flush=True)
	return control

	return EvenCheckpoints()


	def prediction_logits(prediction):
	"""Trainer may return logits alone or a tuple containing logits and auxiliary outputs."""
	return prediction.predictions[0] if isinstance(prediction.predictions, tuple) else prediction.predictions


	def compute_metrics(prediction):
	return classification_metrics(
	prediction.label_ids.tolist(), prediction_logits(prediction).argmax(axis=-1).tolist()
	)


	def softmax_rows(logits):
	import numpy as np

	shifted = logits - logits.max(axis=-1, keepdims=True)
	exp = np.exp(shifted)
	return exp / exp.sum(axis=-1, keepdims=True)


	def prediction_rows(rows, logits):
	"""One record per input with the gold label, predicted label and class probabilities."""
	probabilities = softmax_rows(logits.astype("float64"))
	keep = ("pair_id", "case_id", "domain", "policy_id", "conv_id", "synthetic", "selection_tier",
	"message", "policy", "label")
	return [
	{
	**{key: row.get(key) for key in keep},
	"predicted": LABELS[int(probs.argmax())],
	"probabilities": {name: float(probs[index]) for index, name in enumerate(LABELS)},
	}
	for row, probs in zip(rows, probabilities, strict=True)
	]


	def evaluation_report(records):
	"""Metrics, confusion matrix and the human-readable sklearn report from prediction records.

	`by_source` repeats metrics and confusion matrix for natural pairs (original replies) and edited
	pairs (replies changed by the edit stage) separately. Most training violations are edits, so the
	overall numbers can look good while natural violations are missed; the natural rows are the ones
	that resemble production.
	"""
	from sklearn.metrics import classification_report

	truth = [LABEL2ID[record["label"]] for record in records]
	predicted = [LABEL2ID[record["predicted"]] for record in records]
	by_source = {}
	for source, is_edited in (("natural", False), ("edited", True)):
	subset = [record for record in records if bool(record.get("synthetic")) == is_edited]
	if subset:
	subset_truth = [LABEL2ID[record["label"]] for record in subset]
	subset_predicted = [LABEL2ID[record["predicted"]] for record in subset]
	by_source[source] = {
	"pairs": len(subset),
	"metrics": classification_metrics(subset_truth, subset_predicted),
	"confusion_matrix": confusion_matrix_counts(subset_truth, subset_predicted),
	}
	return {
	"metrics": classification_metrics(truth, predicted),
	"confusion_matrix": confusion_matrix_counts(truth, predicted),
	"by_source": by_source,
	"classification_report": classification_report(
	truth, predicted, labels=list(range(len(LABELS))), target_names=list(LABELS),
	digits=4, zero_division=0,
	),
	}


	def source_metrics_for_logging(split, report):
	"""Flat W&B keys such as validation_natural/violated_recall from evaluation_report's by_source."""
	return {
	f"{split}_{source}/{name}": value
	for source, part in report["by_source"].items()
	for name, value in part["metrics"].items()
	}


	def require_online_wandb(config):
	"""Fail before training if W&B would log offline. The 30k runs logged offline (set by the host's
	environment), so they had no run URLs until synced by hand. Set allow_offline_wandb: true in the
	config to permit it on a host without network access."""
	import os

	mode = os.environ.get("WANDB_MODE", "online").lower()
	if mode in {"offline", "dryrun", "disabled"} and not config.get("allow_offline_wandb", False):
	raise RuntimeError(
	f"WANDB_MODE={mode}: runs would not appear in W&B. Unset WANDB_MODE (and log in with "
	"WANDB_API_KEY), or set allow_offline_wandb: true in the config."
	)


	def run_training(config, model, tokenizer, base_config, collator, resume_from_checkpoint=None):
	"""Train with HF Trainer and W&B logging; adapters saved at every checkpoint fraction."""
	import torch
	import wandb
	from transformers import Trainer, TrainingArguments, set_seed

	from datasets import Dataset

	if not torch.cuda.is_available() or not torch.cuda.is_bf16_supported():
	raise RuntimeError("Training requires a CUDA GPU with BF16 support (configured for A100)")
	require_online_wandb(config)
	splits = validate_inputs(config)
	set_seed(config["seed"])

	def tokenize(batch):
	encoded = tokenize_pairs(tokenizer, batch["message"], batch["policy"], config["max_length"])
	encoded["labels"] = batch["labels"]
	return encoded

	datasets = {
	name: Dataset.from_list(rows).map(tokenize, batched=True, remove_columns=list(rows[0]))
	for name, rows in splits.items()
	if rows
	}
	output = Path(config["output_dir"])
	output.mkdir(parents=True, exist_ok=True)
	(output / "training_config.json").write_text(json.dumps(config, indent=2), encoding="utf-8")
	recipe = input_recipe(config, tokenizer, base_config)
	# transformers 5 folded warmup_ratio into warmup_steps (a float < 1 is a ratio) and dropped
	# group_by_length; 4.57 keeps the original arguments, so behaviour there is unchanged.
	fields = TrainingArguments.__dataclass_fields__
	version_args = {}
	if "warmup_ratio" in fields:
	version_args["warmup_ratio"] = config["warmup_ratio"]
	else:
	version_args["warmup_steps"] = float(config["warmup_ratio"])
	if "group_by_length" in fields:
	version_args["group_by_length"] = config.get("group_by_length", False)
	elif config.get("group_by_length", False):
	raise ValueError("group_by_length is not supported by this transformers version")
	training_args = TrainingArguments(
	**version_args,
	output_dir=str(output),
	run_name=config["run_name"],
	num_train_epochs=config["epochs"],
	learning_rate=config["learning_rate"],
	lr_scheduler_type=config["lr_scheduler_type"],
	weight_decay=config["weight_decay"],
	max_grad_norm=config["max_grad_norm"],
	per_device_train_batch_size=config["batch_size"],
	per_device_eval_batch_size=config["eval_batch_size"],
	gradient_accumulation_steps=config["gradient_accumulation_steps"],
	bf16=True,
	tf32=True,
	gradient_checkpointing=config["gradient_checkpointing"],
	gradient_checkpointing_kwargs={"use_reentrant": False},
	# Evaluation and saving are driven by the checkpoint callback, never by Trainer itself.
	eval_strategy="no",
	save_strategy="no",
	logging_steps=config["logging_steps"],
	report_to=["wandb"],
	seed=config["seed"],
	data_seed=config["seed"],
	optim="adamw_torch",
	dataloader_num_workers=config["dataloader_num_workers"],
	eval_accumulation_steps=8,
	)
	trainer = Trainer(
	model=model,
	args=training_args,
	train_dataset=datasets["train"],
	eval_dataset=datasets["validation"],
	processing_class=tokenizer,
	data_collator=collator,
	compute_metrics=compute_metrics,
	callbacks=[quarter_checkpoint_callback(tokenizer, recipe, output, config["checkpoints_per_run"])],
	)
	# Open the W&B run only after setup succeeds; the finally block closes it on errors too.
	run = wandb.init(project=config["wandb_project"], name=config["run_name"], config=config)
	(output / "wandb_run.json").write_text(json.dumps({"url": run.url, "id": run.id}), encoding="utf-8")
	try:
	torch.cuda.reset_peak_memory_stats()
	started = time.perf_counter()
	trained = trainer.train(resume_from_checkpoint=resume_from_checkpoint)
	torch.cuda.synchronize()
	hardware = {
	"gpu": torch.cuda.get_device_name(),
	"peak_allocated_gib": torch.cuda.max_memory_allocated() / 2**30,
	"peak_reserved_gib": torch.cuda.max_memory_reserved() / 2**30,
	"training_wall_seconds": time.perf_counter() - started,
	}
	(output / "hardware_metrics.json").write_text(json.dumps(hardware, indent=2), encoding="utf-8")
	wandb.log(hardware)
	trainer.save_metrics("train", trained.metrics)
	trainer.save_state()
	save_adapter(model, tokenizer, recipe, output / "final") # same weights as the last checkpoint
	for split in ("validation", "test_ood"):
	if split not in datasets:
	continue
	prediction = trainer.predict(datasets[split], metric_key_prefix=split)
	records = prediction_rows(splits[split], prediction_logits(prediction))
	report = evaluation_report(records)
	write_rows(output / f"{split}_predictions.jsonl", records)
	(output / f"{split}_report.json").write_text(json.dumps(report, indent=2), encoding="utf-8")
	print(report["classification_report"], flush=True)
	for source, part in report["by_source"].items():
	metrics = part["metrics"]
	print(f"{split} {source} ({part['pairs']} pairs): macro F1 {metrics['macro_f1']:.4f}, "
	f"violated precision {metrics['violated_precision']:.4f}, "
	f"violated recall {metrics['violated_recall']:.4f}", flush=True)
	trainer.log_metrics(split, prediction.metrics)
	wandb.log(prediction.metrics)
	wandb.log(source_metrics_for_logging(split, report))
	wandb.log(
	{
	f"{split}_confusion_matrix": wandb.plot.confusion_matrix(
	y_true=[LABEL2ID[r["label"]] for r in records],
	preds=[LABEL2ID[r["predicted"]] for r in records],
	class_names=list(LABELS),
	),
	f"{split}_errors": wandb.Table(
	columns=["message", "policy", "actual", "predicted", "p_violated"],
	data=[
	[r["message"], r["policy"], r["label"], r["predicted"],
	r["probabilities"]["violated"]]
	for r in records if r["label"] != r["predicted"]
	][:500],
	),
	}
	)
	finally:
	wandb.finish()


	def load_trained_classifier(checkpoint, device):
	"""Rebuild base model + saved adapter/head from a checkpoint directory's input_format.json."""
	import torch
	from peft import PeftModel
	from transformers import AutoModelForSequenceClassification, AutoTokenizer

	checkpoint = Path(checkpoint)
	recipe = json.loads((checkpoint / "input_format.json").read_text(encoding="utf-8"))
	tokenizer = AutoTokenizer.from_pretrained(checkpoint)
	dtype = torch.bfloat16 if device == "cuda" and torch.cuda.is_bf16_supported() else torch.float32
	extra = {"reference_compile": False} if recipe["model_type"] == "modernbert" else {"use_cache": False}
	if recipe["model_type"] == "qwen3_5_text":
	# Qwen3.5 checkpoints are vision-language; the classifier uses only the text decoder.
	from training.train_qwen35 import load_text_classifier

	base = load_text_classifier(
	recipe["model_name"], recipe.get("model_revision"), recipe["label2id"], tokenizer.pad_token_id,
	dtype, recipe.get("attention_implementation", "sdpa"),
	)
	model = PeftModel.from_pretrained(base, checkpoint).to(device).eval()
	return model, tokenizer, recipe
	base = AutoModelForSequenceClassification.from_pretrained(
	recipe["model_name"],
	revision=recipe.get("model_revision"),
	num_labels=len(recipe["label2id"]),
	label2id=recipe["label2id"],
	id2label={index: label for label, index in recipe["label2id"].items()},
	pad_token_id=tokenizer.pad_token_id,
	torch_dtype=dtype,
	attn_implementation=recipe.get("attention_implementation", "sdpa"),
	**extra,
	)
	model = PeftModel.from_pretrained(base, checkpoint).to(device).eval()
	return model, tokenizer, recipe


	def predict_logits(model, tokenizer, rows, recipe, batch_size, device):
	"""Batched inference in the saved input format; returns a numpy array of logits."""
	import numpy as np
	import torch

	# Longest-first batches waste the least padding; results are put back in input order.
	order = sorted(range(len(rows)), key=lambda index: -len(rows[index]["message"]) - len(rows[index]["policy"]))
	logits = [None] * len(rows)
	with torch.inference_mode():
	for start in range(0, len(order), batch_size):
	chunk = order[start : start + batch_size]
	batch = tokenize_pairs(
	tokenizer,
	[rows[index]["message"] for index in chunk],
	[rows[index]["policy"] for index in chunk],
	recipe["max_length"],
	recipe["instruction"],
	padding=True,
	return_tensors="pt",
	).to(device)
	output = model(**batch).logits.float().cpu().numpy()
	for index, value in zip(chunk, output, strict=True):
	logits[index] = value
	print(f"predicted {min(start + batch_size, len(order))}/{len(order)}", flush=True)
	return np.stack(logits)