#!/usr/bin/env python3 """Fine-tune the multi-task purpose-deep ModernBERT classifier with MLX.""" from __future__ import annotations import argparse import json import math import random import shutil import sys import time from collections import Counter from pathlib import Path from typing import Any, Iterator, Sequence import numpy as np from deep_contract import ( DEEP_VARIANTS, HEAD_TOKENS, MAX_LENGTH, SCORABLE_HARD_SLICES, TAIL_TOKENS, DeepTargets, DeepVariant, best_mixed_threshold, encode_fixed_shape_numpy, encode_targets, multitask_metrics, validate_deep_records, validate_variant_config, ) from purpose_data import LABELS, DataError, load_jsonl, write_json from train import ( _fit_temperature, choose_confidence_thresholds, expected_calibration_error, ) from train_mlx import _configure_mlx_device, _linear_schedule SCRIPT_DIR = Path(__file__).resolve().parent DEFAULT_DATASET_DIR = SCRIPT_DIR / ".artifacts" / "dataset-v1" DEFAULT_OUTPUT_ROOT = SCRIPT_DIR / "outputs" def _load_mlx(device: str) -> tuple[Any, Any, Any]: try: import mlx.core as mx import mlx.nn as nn import mlx.optimizers as optim except ImportError as exc: raise DataError( "purpose-deep MLX training requires requirements-mlx.txt" ) from exc _configure_mlx_device(mx, device) return mx, nn, optim def _resolve_source(variant: DeepVariant, local_model: Path | None) -> Path: if local_model is not None: source = local_model.expanduser().resolve() if not source.is_dir(): raise DataError(f"{source}: --model must be a local checkpoint directory") return source try: from huggingface_hub import snapshot_download except ImportError as exc: raise DataError("downloading ModernBERT requires huggingface_hub") from exc print( f"resolving {variant.model_id}@{variant.revision} ({variant.parameter_class})", flush=True, ) return Path( snapshot_download( repo_id=variant.model_id, revision=variant.revision, allow_patterns=[ "config.json", "model.safetensors", "tokenizer.json", "tokenizer_config.json", "special_tokens_map.json", ], ) ) def _load_config(source: Path, variant: DeepVariant) -> dict[str, Any]: path = source / "config.json" try: config = json.loads(path.read_text(encoding="utf-8")) except (OSError, json.JSONDecodeError) as exc: raise DataError(f"{path}: cannot load ModernBERT config: {exc}") from exc validate_variant_config(config, variant) return config def _prepare_output(path: Path, source: Path, overwrite: bool) -> None: try: source.resolve().relative_to(path.resolve()) except ValueError: pass else: raise DataError("the input checkpoint must not be inside --output-dir") if path.exists() and any(path.iterdir()): if not overwrite: raise DataError( f"{path}: output is not empty; pass --overwrite-output intentionally" ) shutil.rmtree(path) path.mkdir(parents=True, exist_ok=True) def _encode_records( tokenizer: Any, records: Sequence[dict[str, Any]], *, chunk_size: int = 256, ) -> dict[str, np.ndarray]: chunks: dict[str, list[np.ndarray]] = {} for start in range(0, len(records), chunk_size): encoded = encode_fixed_shape_numpy( tokenizer, [record["prompt"] for record in records[start : start + chunk_size]], ) for key, value in encoded.items(): chunks.setdefault(key, []).append(value) return {key: np.concatenate(values) for key, values in chunks.items()} def _batch_indexes( size: int, batch_size: int, *, permutation: np.ndarray | None = None, ) -> Iterator[np.ndarray]: indexes = permutation if permutation is not None else np.arange(size) for start in range(0, size, batch_size): yield indexes[start : start + batch_size] def _mlx_batch( mx: Any, encoded: dict[str, np.ndarray], targets: DeepTargets, weights: np.ndarray, indexes: np.ndarray, ) -> dict[str, Any]: return { "input_ids": mx.array(encoded["input_ids"][indexes]), "attention_mask": mx.array(encoded["attention_mask"][indexes]), "primary": mx.array(targets.primary[indexes]), "secondary": mx.array(targets.secondary[indexes]), "secondary_mask": mx.array(targets.secondary_mask[indexes]), "mixed": mx.array(targets.mixed[indexes]), "difficulty": mx.array(targets.difficulty[indexes]), "sample_weights": mx.array(weights[indexes]), } def _evaluate( mx: Any, model: Any, encoded: dict[str, np.ndarray], batch_size: int, ) -> dict[str, np.ndarray]: model.eval() collected: dict[str, list[np.ndarray]] = {} for indexes in _batch_indexes(len(encoded["input_ids"]), batch_size): output = model( input_ids=mx.array(encoded["input_ids"][indexes]), attention_mask=mx.array(encoded["attention_mask"][indexes]), ) mx.eval(*output.values()) for key, value in output.items(): collected.setdefault(key, []).append(np.asarray(value)) return {key: np.concatenate(values) for key, values in collected.items()} def _secondary_class_weights(records: Sequence[dict[str, Any]]) -> np.ndarray: counts = Counter( record["secondary"] for record in records if record["secondary"] is not None ) present = [counts[label] for label in LABELS if counts[label]] if not present: raise DataError("purpose-deep needs mixed records with secondary labels") reference = sum(present) / len(present) # Square-root balancing corrects the known skew without letting a five-example # secondary class dominate the shared encoder's primary-purpose gradients. raw = np.asarray( [math.sqrt(reference / max(counts[label], 1)) for label in LABELS], dtype=np.float32, ) return raw / raw.mean() def _sample_weights( records: Sequence[dict[str, Any]], hard_weight: float ) -> np.ndarray: return np.asarray( [hard_weight if record["slice"] in SCORABLE_HARD_SLICES else 1.0 for record in records], dtype=np.float32, ) def _checkpoint_config( source_config: dict[str, Any], variant: DeepVariant, ) -> dict[str, Any]: config = dict(source_config) config.update( { "architectures": ["ModernBertForPurposeClassification"], "id2label": {str(index): label for index, label in enumerate(LABELS)}, "label2id": {label: index for index, label in enumerate(LABELS)}, "num_labels": len(LABELS), "purpose_classifier": { "schemaVersion": 1, "modelVersion": f"purpose-deep-v1-{variant.name}", "trainingBackend": "mlx", "fixedInputShape": [1, MAX_LENGTH], "heads": ["purpose", "secondary", "mixed", "difficulty"], }, } ) return config def _save_checkpoint( mx: Any, model: Any, tokenizer: Any, destination: Path, config: dict[str, Any], ) -> None: from deep_model_mlx import save_weights if destination.exists(): shutil.rmtree(destination) destination.mkdir(parents=True) tokenizer.save_pretrained(destination) (destination / "config.json").write_text( json.dumps(config, indent=2, sort_keys=True) + "\n", encoding="utf-8" ) save_weights(model, destination / "model.safetensors") mx.eval(model.parameters()) def _softmax(values: np.ndarray) -> np.ndarray: shifted = values - values.max(axis=-1, keepdims=True) exponentials = np.exp(shifted) return exponentials / exponentials.sum(axis=-1, keepdims=True) def _calibration( outputs: dict[str, np.ndarray], records: Sequence[dict[str, Any]], args: argparse.Namespace, model_version: str, ) -> tuple[dict[str, Any], dict[str, Any]]: try: import torch except ImportError as exc: raise DataError("final purpose-deep calibration requires PyTorch") from exc targets = encode_targets(records) scorable = np.asarray( [record["slice"] != "vague-eval" for record in records], dtype=np.bool_ ) temperature = _fit_temperature( torch, torch.from_numpy(outputs["purpose_logits"][scorable]), torch.from_numpy(targets.primary[scorable].astype(np.int64)), ) probabilities = _softmax(outputs["purpose_logits"] / temperature) ranked = np.argsort(probabilities, axis=-1) row_indexes = np.arange(len(records)) top = ranked[:, -1] top_probabilities = probabilities[row_indexes, top] margins = top_probabilities - probabilities[row_indexes, ranked[:, -2]] correct = ((top == targets.primary) & scorable).tolist() confidence = choose_confidence_thresholds( top_probabilities.tolist(), margins.tolist(), correct, high_precision=args.high_precision, accepted_precision=args.accepted_precision, ) mixed_mask = targets.secondary_mask secondary_temperature = _fit_temperature( torch, torch.from_numpy(outputs["secondary_logits"][mixed_mask]), torch.from_numpy(targets.secondary[mixed_mask].astype(np.int64)), ) mixed_threshold = best_mixed_threshold( outputs["mixed_logits"][scorable], targets.mixed[scorable] ) calibrated_metrics = multitask_metrics( outputs, records, mixed_threshold=mixed_threshold ) vague = ~scorable score = top_probabilities * (0.5 + 0.5 * margins) vague_low_rate = ( float(np.mean(score[vague] < confidence["medium"]["minimumScore"])) if np.any(vague) else None ) calibration = { "schemaVersion": 1, "modelVersion": model_version, "labels": list(LABELS), "temperature": temperature, "confidence": confidence, "validationECE": expected_calibration_error( top_probabilities.tolist(), correct ), "secondary": { "temperature": secondary_temperature, "labels": list(LABELS), }, "mixed": { "threshold": mixed_threshold, "validationF1": calibrated_metrics["mixed"]["f1"], }, "difficulty": { "activation": "sigmoid", "advisoryOnly": True, }, } return calibration, { "multitask": calibrated_metrics, "vagueLowRate": vague_low_rate, } def train(args: argparse.Namespace) -> dict[str, Any]: mx, nn, optim = _load_mlx(args.device) try: from transformers import AutoTokenizer from deep_model_mlx import ( ModernBertForPurposeClassification, ModernBertPurposeConfig, load_checkpoint_weights, load_pretrained_weights, ) except ImportError as exc: raise DataError( "purpose-deep dependencies are missing; install requirements-base.txt " "and requirements-mlx.txt" ) from exc variant = DEEP_VARIANTS[args.variant] source = _resolve_source(variant, args.resume_from or args.model) source_config = _load_config(source, variant) output_dir = args.output_dir or ( DEFAULT_OUTPUT_ROOT / f"purpose-deep-v1-{variant.name}-mlx" ) _prepare_output(output_dir, source, args.overwrite_output) train_path = args.dataset_dir / "train.jsonl" validation_path = args.dataset_dir / "validation.jsonl" train_records = load_jsonl(train_path) validation_records = load_jsonl(validation_path) validate_deep_records(train_records, str(train_path), training=True) validate_deep_records(validation_records, str(validation_path), training=False) if args.max_train_records: train_records = train_records[: args.max_train_records] if args.max_validation_records: validation_records = validation_records[: args.max_validation_records] tokenizer = AutoTokenizer.from_pretrained(source, local_files_only=True) print("tokenizing fixed 1x512 train and validation splits", flush=True) encoded_train = _encode_records(tokenizer, train_records) encoded_validation = _encode_records(tokenizer, validation_records) train_targets = encode_targets(train_records) validation_targets = encode_targets(validation_records) sample_weights = _sample_weights(train_records, args.hard_weight) secondary_class_weights = _secondary_class_weights(train_records) non_mixed = len(train_records) - int(train_targets.mixed.sum()) mixed_positive_weight = math.sqrt( non_mixed / max(float(train_targets.mixed.sum()), 1.0) ) random.seed(args.seed) np.random.seed(args.seed) mx.random.seed(args.seed) model_config = ModernBertPurposeConfig.from_hugging_face( source_config, gradient_checkpointing=not args.no_gradient_checkpointing ) model = ModernBertForPurposeClassification(model_config) if args.resume_from is not None: load_report = load_checkpoint_weights( model, source / "model.safetensors" ) print( f"resumed purpose-deep tensors={load_report['loaded']} " "including all task heads; optimizer state starts fresh", flush=True, ) else: load_report = load_pretrained_weights(model, source / "model.safetensors") print( f"loaded ModernBERT tensors={load_report['loaded']} " f"ignored_mlm_tensors={load_report['ignored']} " f"fresh_task_tensors={load_report['freshTaskHeads']}", flush=True, ) batch_size = args.batch_size or (4 if variant.name == "base" else 2) eval_batch_size = args.eval_batch_size or (8 if variant.name == "base" else 4) learning_rate = args.learning_rate or ( 2e-5 if variant.name == "base" else 1e-5 ) steps_per_epoch = math.ceil(len(train_records) / batch_size) total_steps = steps_per_epoch * args.epochs optimizer = optim.AdamW( learning_rate=_linear_schedule( mx, learning_rate, total_steps, round(total_steps * args.warmup_ratio), ), weight_decay=args.weight_decay, bias_correction=True, ) class_weights_mx = mx.array(secondary_class_weights) def loss_function( input_ids: Any, attention_mask: Any, primary: Any, secondary: Any, secondary_mask: Any, mixed: Any, difficulty: Any, weights: Any, ) -> tuple[Any, Any, Any, Any, Any]: output = model(input_ids=input_ids, attention_mask=attention_mask) primary_per_record = nn.losses.cross_entropy( output["purpose_logits"], primary, label_smoothing=args.label_smoothing, reduction="none", ) primary_loss = mx.sum(primary_per_record * weights) / mx.sum(weights) secondary_per_record = nn.losses.cross_entropy( output["secondary_logits"], secondary, label_smoothing=args.label_smoothing, reduction="none", ) secondary_weights = ( weights * secondary_mask.astype(weights.dtype) * class_weights_mx[secondary] ) secondary_loss = mx.sum(secondary_per_record * secondary_weights) / mx.maximum( mx.sum(secondary_weights), 1.0 ) mixed_per_record = nn.losses.binary_cross_entropy( output["mixed_logits"], mixed, reduction="none" ) mixed_balance = mx.where(mixed > 0.5, mixed_positive_weight, 1.0) mixed_loss = mx.sum(mixed_per_record * mixed_balance * weights) / mx.sum( mixed_balance * weights ) difficulty_per_record = nn.losses.smooth_l1_loss( output["difficulty"], difficulty, beta=0.1, reduction="none" ) difficulty_loss = mx.sum(difficulty_per_record * weights) / mx.sum(weights) total = ( primary_loss + args.secondary_loss_weight * secondary_loss + args.mixed_loss_weight * mixed_loss + args.difficulty_loss_weight * difficulty_loss ) return total, primary_loss, secondary_loss, mixed_loss, difficulty_loss loss_and_grad = nn.value_and_grad(model, loss_function) rng = np.random.default_rng(args.seed) checkpoint_config = _checkpoint_config(source_config, variant) best_dir = output_dir / "model" best_score = float("-inf") best_metrics: dict[str, Any] | None = None epochs_without_improvement = 0 stopped_early = False history: list[dict[str, Any]] = [] started = time.perf_counter() for epoch in range(1, args.epochs + 1): epoch_started = time.perf_counter() model.train() running = np.zeros(5, dtype=np.float64) permutation = rng.permutation(len(train_records)) for step, indexes in enumerate( _batch_indexes( len(train_records), batch_size, permutation=permutation ), 1, ): batch = _mlx_batch( mx, encoded_train, train_targets, sample_weights, indexes ) losses, gradients = loss_and_grad( batch["input_ids"], batch["attention_mask"], batch["primary"], batch["secondary"], batch["secondary_mask"], batch["mixed"], batch["difficulty"], batch["sample_weights"], ) gradients, _ = optim.clip_grad_norm(gradients, args.max_grad_norm) optimizer.update(model, gradients) mx.eval(model.parameters(), optimizer.state, *losses) running += np.asarray([float(value.item()) for value in losses]) if args.progress_steps and ( step % args.progress_steps == 0 or step == steps_per_epoch ): mean = running / step print( f"epoch {epoch} step {step}/{steps_per_epoch} " f"loss={mean[0]:.4f} primary={mean[1]:.4f} " f"secondary={mean[2]:.4f} mixed={mean[3]:.4f} " f"difficulty={mean[4]:.4f} " f"elapsed={time.perf_counter() - epoch_started:.1f}s", flush=True, ) outputs = _evaluate( mx, model, encoded_validation, eval_batch_size ) metrics = multitask_metrics(outputs, validation_records) metrics["epoch"] = epoch metrics["meanTrainingLoss"] = (running / steps_per_epoch).tolist() history.append(metrics) score = float(metrics["selectionScore"]) secondary_macro = ( metrics["secondary"]["macroRecall"] if metrics["secondary"] is not None else 0.0 ) print( f"epoch {epoch}: primary_accuracy={metrics['primary']['accuracy']:.4%} " f"hard_accuracy={metrics['primaryHardSlice']['accuracy']:.4%} " f"secondary_macro_recall={secondary_macro:.4%} " f"mixed_f1={metrics['mixed']['f1']:.4%} " f"difficulty_mae={metrics['difficulty']['mae']:.4f} " f"selection_score={score:.4%}", flush=True, ) improvement = score - best_score if improvement > args.minimum_improvement: best_score = score best_metrics = metrics epochs_without_improvement = 0 _save_checkpoint( mx, model, tokenizer, best_dir, checkpoint_config ) write_json( output_dir / "training-state.json", { "bestEpoch": epoch, "bestSelectionScore": best_score, "elapsedSeconds": time.perf_counter() - started, "complete": False, }, ) else: epochs_without_improvement += 1 if epochs_without_improvement >= args.early_stopping_patience: stopped_early = True print( f"early stopping after epoch {epoch}: no hard-aware " f"selection improvement greater than " f"{args.minimum_improvement:.4%} for " f"{args.early_stopping_patience} epoch(s)", flush=True, ) break if best_metrics is None: raise DataError("purpose-deep training did not produce a checkpoint") # Release the optimizer graph before opening the selected checkpoint; base and # especially large should never hold two full optimizer states at calibration time. del optimizer, loss_and_grad, model mx.clear_cache() selected_model = ModernBertForPurposeClassification(model_config) selected_model.load_weights(str(best_dir / "model.safetensors"), strict=True) selected_outputs = _evaluate( mx, selected_model, encoded_validation, eval_batch_size ) model_version = f"purpose-deep-v1-{variant.name}" calibration, calibrated = _calibration( selected_outputs, validation_records, args, model_version ) metrics = { "modelVersion": model_version, "variant": variant.name, "baseModel": variant.model_id, "baseModelRevision": variant.revision, "resumedFrom": str(source) if args.resume_from is not None else None, "parameterClass": variant.parameter_class, "trainingBackend": "mlx", "device": args.device, "fixedInputShape": [1, MAX_LENGTH], "truncation": { "strategy": "head-tail-pair", "headTokens": HEAD_TOKENS, "tailTokens": TAIL_TOKENS, }, "trainingSeconds": time.perf_counter() - started, "trainRecords": len(train_records), "validationRecords": len(validation_records), "mixedTrainRecords": int(train_targets.mixed.sum()), "mixedValidationRecords": int(validation_targets.mixed.sum()), "hardTrainingWeight": args.hard_weight, "lossWeights": { "purpose": 1.0, "secondary": args.secondary_loss_weight, "mixed": args.mixed_loss_weight, "difficulty": args.difficulty_loss_weight, }, "secondaryClassWeights": { label: float(secondary_class_weights[index]) for index, label in enumerate(LABELS) }, "mixedPositiveWeight": mixed_positive_weight, "gradientCheckpointing": not args.no_gradient_checkpointing, "batchSize": batch_size, "learningRate": learning_rate, "bestValidationSelectionScore": best_score, "bestValidation": best_metrics, "selectedValidation": calibrated, "epochsCompleted": len(history), "stoppedEarly": stopped_early, "history": history, "calibration": calibration, } write_json(output_dir / "calibration.json", calibration) write_json(output_dir / "metrics.json", metrics) write_json( output_dir / "training-config.json", { key: str(value) if isinstance(value, Path) else value for key, value in vars(args).items() }, ) write_json( output_dir / "training-state.json", { "bestEpoch": int(best_metrics["epoch"]), "bestSelectionScore": best_score, "elapsedSeconds": metrics["trainingSeconds"], "complete": True, }, ) return metrics def build_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--variant", choices=tuple(DEEP_VARIANTS), default="base") source_group = parser.add_mutually_exclusive_group() source_group.add_argument( "--model", type=Path, help="local pinned ModernBERT checkpoint (default: download the pinned revision)", ) source_group.add_argument( "--resume-from", type=Path, help=( "selected purpose-deep model directory to continue from; restores " "the backbone and all four task heads with a fresh optimizer" ), ) parser.add_argument("--dataset-dir", type=Path, default=DEFAULT_DATASET_DIR) parser.add_argument("--output-dir", type=Path) parser.add_argument( "--device", choices=("metal", "cpu"), default="metal", help="MLX execution device (Metal by default; CPU is diagnostic only)", ) parser.add_argument("--seed", type=int, default=20260731) parser.add_argument("--epochs", type=int, default=3) parser.add_argument("--batch-size", type=int) parser.add_argument("--eval-batch-size", type=int) parser.add_argument("--learning-rate", type=float) parser.add_argument("--weight-decay", type=float, default=0.01) parser.add_argument("--warmup-ratio", type=float, default=0.1) parser.add_argument("--max-grad-norm", type=float, default=1.0) parser.add_argument("--label-smoothing", type=float, default=0.05) parser.add_argument("--hard-weight", type=float, default=2.0) parser.add_argument("--secondary-loss-weight", type=float, default=0.25) parser.add_argument("--mixed-loss-weight", type=float, default=0.25) parser.add_argument("--difficulty-loss-weight", type=float, default=0.10) parser.add_argument("--progress-steps", type=int, default=25) parser.add_argument("--early-stopping-patience", type=int, default=1) parser.add_argument("--minimum-improvement", type=float, default=0.0005) parser.add_argument("--high-precision", type=float, default=0.98) parser.add_argument("--accepted-precision", type=float, default=0.95) parser.add_argument("--no-gradient-checkpointing", action="store_true") parser.add_argument("--max-train-records", type=int) parser.add_argument("--max-validation-records", type=int) parser.add_argument("--overwrite-output", action="store_true") return parser def _positive(parser: argparse.ArgumentParser, name: str, value: Any) -> None: if value is not None and value <= 0: parser.error(f"--{name.replace('_', '-')} must be positive") def main(argv: Sequence[str] | None = None) -> int: parser = build_parser() args = parser.parse_args(argv) for name in ( "epochs", "batch_size", "eval_batch_size", "learning_rate", "max_grad_norm", "hard_weight", "early_stopping_patience", ): _positive(parser, name, getattr(args, name)) if args.progress_steps < 0: parser.error("--progress-steps must be non-negative") if not 0 <= args.warmup_ratio < 1: parser.error("--warmup-ratio must be in [0, 1)") if not 0 <= args.label_smoothing < 1: parser.error("--label-smoothing must be in [0, 1)") for name in ( "secondary_loss_weight", "mixed_loss_weight", "difficulty_loss_weight", ): if getattr(args, name) < 0: parser.error(f"--{name.replace('_', '-')} must be non-negative") if not 0 < args.accepted_precision <= args.high_precision <= 1: parser.error( "confidence precision targets must satisfy 0 < accepted <= high <= 1" ) try: metrics = train(args) except (DataError, OSError, RuntimeError, ValueError) as exc: print(f"error: {exc}", file=sys.stderr) return 1 selected = metrics["selectedValidation"]["multitask"] print( f"selected validation: primary={selected['primary']['accuracy']:.4%} " f"hard={selected['primaryHardSlice']['accuracy']:.4%} " f"mixed_f1={selected['mixed']['f1']:.4%}", flush=True, ) return 0 if __name__ == "__main__": raise SystemExit(main())