920 lines
33 KiB
Python
920 lines
33 KiB
Python
#!/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, _teacher_cache
|
|
|
|
|
|
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 _distillation_loss(
|
|
mx: Any,
|
|
student_logits: Any,
|
|
teacher_logits: Any,
|
|
*,
|
|
temperature: float,
|
|
weights: Any,
|
|
) -> Any:
|
|
"""Return weighted teacher-to-student KL loss for one MLX batch."""
|
|
|
|
student_log_probabilities = (
|
|
student_logits / temperature
|
|
- mx.logsumexp(student_logits / temperature, axis=-1, keepdims=True)
|
|
)
|
|
teacher_probabilities = mx.softmax(
|
|
teacher_logits / temperature,
|
|
axis=-1,
|
|
)
|
|
teacher_log_probabilities = mx.log(
|
|
mx.maximum(teacher_probabilities, 1e-12)
|
|
)
|
|
per_record = (
|
|
mx.sum(
|
|
teacher_probabilities
|
|
* (teacher_log_probabilities - student_log_probabilities),
|
|
axis=-1,
|
|
)
|
|
* temperature
|
|
* temperature
|
|
)
|
|
return mx.sum(per_record * weights) / mx.sum(weights)
|
|
|
|
|
|
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)
|
|
validation_scorable = np.asarray(
|
|
[record["slice"] != "vague-eval" for record in validation_records],
|
|
dtype=np.bool_,
|
|
)
|
|
sample_weights = _sample_weights(train_records, args.hard_weight)
|
|
secondary_class_weights = _secondary_class_weights(train_records)
|
|
teacher_train_logits = None
|
|
teacher_validation_logits = None
|
|
if args.distillation_cache is not None:
|
|
teacher_train_logits, teacher_validation_logits = _teacher_cache(
|
|
args.distillation_cache.expanduser(),
|
|
train_records,
|
|
validation_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,
|
|
teacher_logits: Any | None,
|
|
) -> tuple[Any, 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_label_loss = mx.sum(primary_per_record * weights) / mx.sum(weights)
|
|
primary_distillation_loss = mx.zeros_like(primary_label_loss)
|
|
if teacher_logits is not None:
|
|
primary_distillation_loss = _distillation_loss(
|
|
mx,
|
|
output["purpose_logits"],
|
|
teacher_logits,
|
|
temperature=args.distillation_temperature,
|
|
weights=weights,
|
|
)
|
|
primary_loss = (
|
|
(1.0 - args.distillation_weight) * primary_label_loss
|
|
+ args.distillation_weight * primary_distillation_loss
|
|
)
|
|
|
|
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_label_loss,
|
|
secondary_loss,
|
|
mixed_loss,
|
|
difficulty_loss,
|
|
primary_distillation_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"
|
|
epochs_without_improvement = 0
|
|
stopped_early = False
|
|
history: list[dict[str, Any]] = []
|
|
started = time.perf_counter()
|
|
|
|
initial_outputs = _evaluate(
|
|
mx, model, encoded_validation, eval_batch_size
|
|
)
|
|
initial_metrics = multitask_metrics(initial_outputs, validation_records)
|
|
if teacher_validation_logits is not None:
|
|
initial_metrics["teacherAgreement"] = float(
|
|
np.mean(
|
|
initial_outputs["purpose_logits"][validation_scorable].argmax(
|
|
axis=-1
|
|
)
|
|
== teacher_validation_logits[validation_scorable].argmax(axis=-1)
|
|
)
|
|
)
|
|
initial_metrics["epoch"] = 0
|
|
best_score = float(initial_metrics["selectionScore"])
|
|
best_metrics: dict[str, Any] = initial_metrics
|
|
_save_checkpoint(mx, model, tokenizer, best_dir, checkpoint_config)
|
|
write_json(
|
|
output_dir / "training-state.json",
|
|
{
|
|
"bestEpoch": 0,
|
|
"bestSelectionScore": best_score,
|
|
"elapsedSeconds": time.perf_counter() - started,
|
|
"complete": False,
|
|
},
|
|
)
|
|
print(
|
|
f"epoch 0: primary_accuracy={initial_metrics['primary']['accuracy']:.4%} "
|
|
f"hard_accuracy={initial_metrics['primaryHardSlice']['accuracy']:.4%} "
|
|
f"mixed_f1={initial_metrics['mixed']['f1']:.4%} "
|
|
f"selection_score={best_score:.4%}"
|
|
+ (
|
|
f" teacher_agreement={initial_metrics['teacherAgreement']:.4%}"
|
|
if "teacherAgreement" in initial_metrics
|
|
else ""
|
|
),
|
|
flush=True,
|
|
)
|
|
|
|
for epoch in range(1, args.epochs + 1):
|
|
epoch_started = time.perf_counter()
|
|
model.train()
|
|
running = np.zeros(6, 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
|
|
)
|
|
teacher_logits = (
|
|
mx.array(teacher_train_logits[indexes])
|
|
if teacher_train_logits is not None
|
|
else None
|
|
)
|
|
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"],
|
|
teacher_logits,
|
|
)
|
|
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" distillation={mean[5]:.4f}"
|
|
if teacher_train_logits is not None
|
|
else ""
|
|
)
|
|
+ " "
|
|
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)
|
|
if teacher_validation_logits is not None:
|
|
metrics["teacherAgreement"] = float(
|
|
np.mean(
|
|
outputs["purpose_logits"][validation_scorable].argmax(axis=-1)
|
|
== teacher_validation_logits[validation_scorable].argmax(axis=-1)
|
|
)
|
|
)
|
|
metrics["epoch"] = epoch
|
|
metrics["meanTrainingLoss"] = (running / steps_per_epoch).tolist()
|
|
history.append(metrics)
|
|
score = float(metrics["selectionScore"])
|
|
secondary_macro = (
|
|
metrics["secondary"]["supportedMacroRecall"]
|
|
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_supported_macro_recall={secondary_macro:.4%} "
|
|
f"mixed_f1={metrics['mixed']['f1']:.4%} "
|
|
f"difficulty_mae={metrics['difficulty']['mae']:.4f} "
|
|
f"selection_score={score:.4%}"
|
|
+ (
|
|
f" teacher_agreement={metrics['teacherAgreement']:.4%}"
|
|
if "teacherAgreement" in metrics
|
|
else ""
|
|
),
|
|
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
|
|
|
|
# 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,
|
|
},
|
|
"distillation": {
|
|
"cache": (
|
|
str(args.distillation_cache)
|
|
if args.distillation_cache is not None
|
|
else None
|
|
),
|
|
"weight": args.distillation_weight,
|
|
"temperature": args.distillation_temperature,
|
|
},
|
|
"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,
|
|
"initialValidation": initial_metrics,
|
|
"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("--distillation-cache", type=Path)
|
|
parser.add_argument("--distillation-weight", type=float, default=0.0)
|
|
parser.add_argument("--distillation-temperature", type=float, default=2.0)
|
|
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.distillation_weight <= 1:
|
|
parser.error("--distillation-weight must be in [0, 1]")
|
|
if args.distillation_temperature <= 0:
|
|
parser.error("--distillation-temperature must be positive")
|
|
if (args.distillation_cache is None) != (args.distillation_weight == 0):
|
|
parser.error(
|
|
"--distillation-cache and a positive --distillation-weight "
|
|
"must be supplied together"
|
|
)
|
|
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())
|