1288 lines
46 KiB
Python
1288 lines
46 KiB
Python
#!/usr/bin/env python3
|
|
"""Fine-tune the multi-task purpose-deep ModernBERT classifier with MLX."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import hashlib
|
|
import json
|
|
import math
|
|
import os
|
|
import random
|
|
import shutil
|
|
import shlex
|
|
import signal
|
|
import sys
|
|
import tempfile
|
|
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"
|
|
RESUME_SCHEMA_VERSION = 1
|
|
PATH_ARGUMENTS = {"dataset_dir", "distillation_cache"}
|
|
|
|
|
|
class TrainingPaused(Exception):
|
|
"""Raised after a signal-requested training checkpoint is durable."""
|
|
|
|
def __init__(self, checkpoint: Path, signum: int) -> None:
|
|
super().__init__(str(checkpoint))
|
|
self.checkpoint = checkpoint
|
|
self.signum = signum
|
|
|
|
|
|
class _ShutdownController:
|
|
def __init__(self) -> None:
|
|
self.signum: int | None = None
|
|
self._previous: dict[int, Any] = {}
|
|
|
|
def _handle(self, signum: int, _frame: Any) -> None:
|
|
if self.signum is not None:
|
|
raise KeyboardInterrupt
|
|
self.signum = signum
|
|
|
|
def install(self) -> None:
|
|
for signum in (signal.SIGINT, signal.SIGTERM):
|
|
self._previous[signum] = signal.getsignal(signum)
|
|
signal.signal(signum, self._handle)
|
|
|
|
def restore(self) -> None:
|
|
for signum, handler in self._previous.items():
|
|
signal.signal(signum, handler)
|
|
self._previous.clear()
|
|
|
|
|
|
def _sha256(path: Path) -> str:
|
|
digest = hashlib.sha256()
|
|
with path.open("rb") as source:
|
|
for chunk in iter(lambda: source.read(1024 * 1024), b""):
|
|
digest.update(chunk)
|
|
return digest.hexdigest()
|
|
|
|
|
|
def _read_resume_state(checkpoint: Path) -> dict[str, Any]:
|
|
state_path = checkpoint / "resume-state.json"
|
|
try:
|
|
state = json.loads(state_path.read_text(encoding="utf-8"))
|
|
except (OSError, json.JSONDecodeError) as exc:
|
|
raise DataError(f"{state_path}: cannot load training resume state: {exc}") from exc
|
|
if state.get("schemaVersion") != RESUME_SCHEMA_VERSION:
|
|
raise DataError(f"{state_path}: unsupported training resume schema")
|
|
if state.get("status") != "paused":
|
|
raise DataError(f"{state_path}: checkpoint is not paused training state")
|
|
return state
|
|
|
|
|
|
def _serialized_resume_arguments(args: argparse.Namespace) -> dict[str, Any]:
|
|
excluded = {
|
|
"model",
|
|
"resume_from",
|
|
"resume_training",
|
|
"output_dir",
|
|
"overwrite_output",
|
|
}
|
|
return {
|
|
key: (
|
|
str(value.expanduser().resolve())
|
|
if isinstance(value, Path)
|
|
else value
|
|
)
|
|
for key, value in vars(args).items()
|
|
if key not in excluded
|
|
}
|
|
|
|
|
|
def _restore_resume_arguments(
|
|
args: argparse.Namespace, checkpoint: Path, state: dict[str, Any]
|
|
) -> None:
|
|
saved = state.get("arguments")
|
|
if not isinstance(saved, dict):
|
|
raise DataError("training resume state has no saved arguments")
|
|
for key, value in saved.items():
|
|
if not hasattr(args, key):
|
|
raise DataError(f"training resume state has unknown argument {key!r}")
|
|
setattr(args, key, Path(value) if key in PATH_ARGUMENTS and value else value)
|
|
args.model = None
|
|
args.resume_from = None
|
|
args.resume_training = checkpoint
|
|
args.output_dir = checkpoint.parent
|
|
args.overwrite_output = False
|
|
|
|
|
|
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 _save_training_resume(
|
|
mx: Any,
|
|
model: Any,
|
|
optimizer: Any,
|
|
tokenizer: Any,
|
|
output_dir: Path,
|
|
checkpoint_config: dict[str, Any],
|
|
state: dict[str, Any],
|
|
) -> Path:
|
|
"""Atomically save the current model, optimizer, and loop cursor."""
|
|
|
|
from mlx.utils import tree_flatten
|
|
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
temporary = Path(
|
|
tempfile.mkdtemp(prefix=".training-resume-", dir=output_dir)
|
|
)
|
|
destination = output_dir / "resume"
|
|
backup = output_dir / ".training-resume-backup"
|
|
try:
|
|
_save_checkpoint(mx, model, tokenizer, temporary, checkpoint_config)
|
|
mx.eval(optimizer.state)
|
|
optimizer_state = tree_flatten(optimizer.state, destination={})
|
|
if not optimizer_state:
|
|
raise DataError("optimizer state is empty; refusing an incomplete resume")
|
|
mx.save_safetensors(
|
|
str(temporary / "optimizer.safetensors"), optimizer_state
|
|
)
|
|
write_json(temporary / "resume-state.json", state)
|
|
if backup.exists():
|
|
shutil.rmtree(backup)
|
|
if destination.exists():
|
|
os.replace(destination, backup)
|
|
os.replace(temporary, destination)
|
|
if backup.exists():
|
|
shutil.rmtree(backup)
|
|
except BaseException:
|
|
if temporary.exists():
|
|
shutil.rmtree(temporary)
|
|
if not destination.exists() and backup.exists():
|
|
os.replace(backup, destination)
|
|
raise
|
|
return destination
|
|
|
|
|
|
def _load_optimizer_state(mx: Any, optimizer: Any, checkpoint: Path) -> None:
|
|
from mlx.utils import tree_unflatten
|
|
|
|
path = checkpoint / "optimizer.safetensors"
|
|
if not path.is_file():
|
|
raise DataError(f"{path}: optimizer resume state is missing")
|
|
optimizer.state = tree_unflatten(mx.load(str(path)))
|
|
mx.eval(optimizer.state)
|
|
|
|
|
|
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]
|
|
resume_state = (
|
|
_read_resume_state(args.resume_training)
|
|
if args.resume_training is not None
|
|
else None
|
|
)
|
|
source = _resolve_source(
|
|
variant, args.resume_training or args.resume_from or args.model
|
|
)
|
|
source_config = _load_config(source, variant)
|
|
if args.resume_training is not None:
|
|
output_dir = args.resume_training.parent
|
|
if args.output_dir is not None and args.output_dir.resolve() != output_dir:
|
|
raise DataError("--resume-training must use its original output directory")
|
|
if not (output_dir / "model" / "model.safetensors").is_file():
|
|
raise DataError(f"{output_dir}: selected model checkpoint is missing")
|
|
else:
|
|
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]
|
|
|
|
input_hashes = {
|
|
"train": _sha256(train_path),
|
|
"validation": _sha256(validation_path),
|
|
"distillationCache": (
|
|
_sha256(args.distillation_cache.expanduser())
|
|
if args.distillation_cache is not None
|
|
else None
|
|
),
|
|
}
|
|
if resume_state is not None and resume_state.get("inputHashes") != input_hashes:
|
|
raise DataError(
|
|
"training inputs changed after the pause; refusing a non-deterministic resume"
|
|
)
|
|
|
|
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_training is not None or 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"
|
|
+ (
|
|
" and exact optimizer/loop state"
|
|
if args.resume_training is not None
|
|
else "; 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,
|
|
)
|
|
optimizer.init(model.trainable_parameters())
|
|
mx.eval(optimizer.state)
|
|
if args.resume_training is not None:
|
|
_load_optimizer_state(mx, optimizer, source)
|
|
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"
|
|
stopped_early = False
|
|
write_json(
|
|
output_dir / "training-config.json",
|
|
{
|
|
key: str(value) if isinstance(value, Path) else value
|
|
for key, value in vars(args).items()
|
|
},
|
|
)
|
|
|
|
if resume_state is not None:
|
|
try:
|
|
initial_metrics = resume_state["initialMetrics"]
|
|
best_score = float(resume_state["bestScore"])
|
|
best_metrics = resume_state["bestMetrics"]
|
|
epochs_without_improvement = int(
|
|
resume_state["epochsWithoutImprovement"]
|
|
)
|
|
history = list(resume_state["history"])
|
|
start_epoch = int(resume_state["epoch"])
|
|
resume_next_step = int(resume_state["nextStep"])
|
|
resume_permutation = resume_state.get("permutation")
|
|
resume_running = np.asarray(
|
|
resume_state["runningLosses"], dtype=np.float64
|
|
)
|
|
resume_epoch_elapsed = float(resume_state["epochElapsedSeconds"])
|
|
rng.bit_generator.state = resume_state["numpyRngState"]
|
|
started = time.perf_counter() - float(resume_state["elapsedSeconds"])
|
|
except (KeyError, TypeError, ValueError) as exc:
|
|
raise DataError(f"invalid training loop resume state: {exc}") from exc
|
|
if not 1 <= resume_next_step <= steps_per_epoch + 1:
|
|
raise DataError("training resume step is outside the epoch")
|
|
if resume_running.shape != (6,):
|
|
raise DataError("training resume loss accumulator has the wrong shape")
|
|
print(
|
|
f"continuing epoch {start_epoch} at step "
|
|
f"{resume_next_step}/{steps_per_epoch} after "
|
|
f"{resume_state['elapsedSeconds']:.1f}s of saved training",
|
|
flush=True,
|
|
)
|
|
else:
|
|
epochs_without_improvement = 0
|
|
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,
|
|
)
|
|
start_epoch = 1
|
|
resume_next_step = 1
|
|
resume_permutation = None
|
|
resume_running = np.zeros(6, dtype=np.float64)
|
|
resume_epoch_elapsed = 0.0
|
|
|
|
shutdown = _ShutdownController()
|
|
|
|
def pause_training(
|
|
epoch: int,
|
|
next_step: int,
|
|
permutation: np.ndarray | None,
|
|
running: np.ndarray,
|
|
epoch_elapsed: float,
|
|
) -> None:
|
|
signum = shutdown.signum or signal.SIGINT
|
|
print(
|
|
f"shutdown requested; saving exact training state after epoch {epoch} "
|
|
f"step {max(next_step - 1, 0)}",
|
|
flush=True,
|
|
)
|
|
state = {
|
|
"schemaVersion": RESUME_SCHEMA_VERSION,
|
|
"status": "paused",
|
|
"signal": signal.Signals(signum).name,
|
|
"arguments": _serialized_resume_arguments(args),
|
|
"inputHashes": input_hashes,
|
|
"epoch": epoch,
|
|
"nextStep": next_step,
|
|
"permutation": permutation.tolist() if permutation is not None else None,
|
|
"runningLosses": running.tolist(),
|
|
"epochElapsedSeconds": epoch_elapsed,
|
|
"elapsedSeconds": time.perf_counter() - started,
|
|
"numpyRngState": rng.bit_generator.state,
|
|
"initialMetrics": initial_metrics,
|
|
"bestScore": best_score,
|
|
"bestMetrics": best_metrics,
|
|
"epochsWithoutImprovement": epochs_without_improvement,
|
|
"history": history,
|
|
}
|
|
checkpoint = _save_training_resume(
|
|
mx,
|
|
model,
|
|
optimizer,
|
|
tokenizer,
|
|
output_dir,
|
|
checkpoint_config,
|
|
state,
|
|
)
|
|
write_json(
|
|
output_dir / "training-state.json",
|
|
{
|
|
"bestEpoch": int(best_metrics["epoch"]),
|
|
"bestSelectionScore": best_score,
|
|
"elapsedSeconds": state["elapsedSeconds"],
|
|
"complete": False,
|
|
"paused": True,
|
|
"resumeCheckpoint": str(checkpoint),
|
|
},
|
|
)
|
|
print(
|
|
"training paused safely; resume with:\n"
|
|
f" {shlex.quote(sys.executable)} -u "
|
|
f"{shlex.quote(str(Path(__file__).resolve()))} "
|
|
f"--resume-training {shlex.quote(str(checkpoint))}",
|
|
flush=True,
|
|
)
|
|
raise TrainingPaused(checkpoint, signum)
|
|
|
|
shutdown.install()
|
|
try:
|
|
for epoch in range(start_epoch, args.epochs + 1):
|
|
if resume_state is not None and epoch == start_epoch:
|
|
permutation = (
|
|
np.asarray(resume_permutation, dtype=np.int64)
|
|
if resume_permutation is not None
|
|
else rng.permutation(len(train_records))
|
|
)
|
|
running = resume_running.copy()
|
|
first_step = resume_next_step
|
|
epoch_started = time.perf_counter() - resume_epoch_elapsed
|
|
else:
|
|
permutation = rng.permutation(len(train_records))
|
|
running = np.zeros(6, dtype=np.float64)
|
|
first_step = 1
|
|
epoch_started = time.perf_counter()
|
|
if permutation.shape != (len(train_records),):
|
|
raise DataError("training resume permutation has the wrong shape")
|
|
if shutdown.signum is not None:
|
|
pause_training(
|
|
epoch,
|
|
first_step,
|
|
permutation,
|
|
running,
|
|
time.perf_counter() - epoch_started,
|
|
)
|
|
model.train()
|
|
for step, indexes in enumerate(
|
|
_batch_indexes(
|
|
len(train_records), batch_size, permutation=permutation
|
|
),
|
|
1,
|
|
):
|
|
if step < first_step:
|
|
continue
|
|
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,
|
|
)
|
|
if shutdown.signum is not None:
|
|
pause_training(
|
|
epoch,
|
|
step + 1,
|
|
permutation,
|
|
running,
|
|
time.perf_counter() - epoch_started,
|
|
)
|
|
|
|
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
|
|
resume_state = None
|
|
if shutdown.signum is not None:
|
|
pause_training(
|
|
epoch + 1,
|
|
1,
|
|
None,
|
|
np.zeros(6, dtype=np.float64),
|
|
0.0,
|
|
)
|
|
finally:
|
|
shutdown.restore()
|
|
|
|
# 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 or args.resume_training is not None
|
|
else None
|
|
),
|
|
"exactTrainingResume": args.resume_training is not 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,
|
|
},
|
|
)
|
|
for stale_resume in (
|
|
output_dir / "resume",
|
|
output_dir / ".training-resume-backup",
|
|
):
|
|
if stale_resume.exists():
|
|
shutil.rmtree(stale_resume)
|
|
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"
|
|
),
|
|
)
|
|
source_group.add_argument(
|
|
"--resume-training",
|
|
type=Path,
|
|
help=(
|
|
"exact signal-created resume checkpoint; restores the saved arguments, "
|
|
"model, optimizer, shuffle order, and next batch"
|
|
),
|
|
)
|
|
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)
|
|
if args.resume_training is not None:
|
|
checkpoint = args.resume_training.expanduser().resolve()
|
|
try:
|
|
resume_state = _read_resume_state(checkpoint)
|
|
_restore_resume_arguments(args, checkpoint, resume_state)
|
|
except DataError as exc:
|
|
parser.error(str(exc))
|
|
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 TrainingPaused as paused:
|
|
return 128 + paused.signum
|
|
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())
|