"""Deterministic, dependency-light PyTorch training loop.""" from __future__ import annotations import json import os import random from collections.abc import Iterable, Mapping from dataclasses import asdict, dataclass from pathlib import Path from typing import Any import torch from torch import Tensor, nn from turn_detection.models.common import TurnDetectionOutput from turn_detection.training.losses import MultiTaskLossConfig, MultiTaskTurnLoss from turn_detection.training.metrics import ( binary_classification_metrics, threshold_at_max_fpr, ) def seed_everything(seed: int, deterministic: bool = True) -> None: """Seed Python/PyTorch and select deterministic kernels when requested.""" random.seed(seed) try: import numpy as np except ImportError: np = None if np is not None: np.random.seed(seed) torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) if deterministic: torch.use_deterministic_algorithms(True, warn_only=True) if torch.backends.cudnn.is_available(): torch.backends.cudnn.benchmark = False torch.backends.cudnn.deterministic = True @dataclass(frozen=True) class TrainerConfig: seed: int = 17 epochs: int = 20 learning_rate: float = 3e-4 weight_decay: float = 1e-3 gradient_accumulation_steps: int = 1 max_gradient_norm: float = 1.0 mixed_precision: bool = True deterministic: bool = True device: str = "auto" early_stopping_patience: int = 5 min_improvement: float = 1e-4 false_positive_rate_budget: float = 0.02 selection_metric: str = "constrained_recall" @classmethod def from_mapping(cls, values: Mapping[str, Any]) -> TrainerConfig: known = {field.name for field in cls.__dataclass_fields__.values()} return cls(**{k: v for k, v in values.items() if k in known}) def _resolve_device(requested: str) -> torch.device: if requested != "auto": return torch.device(requested) if torch.cuda.is_available(): return torch.device("cuda") mps = getattr(torch.backends, "mps", None) if mps is not None and mps.is_available(): return torch.device("mps") return torch.device("cpu") def _extract_batch(batch: Mapping[str, Any], device: torch.device) -> tuple[Tensor, ...]: features = batch.get("log_mel", batch.get("input_features", batch.get("features"))) if features is None: raise KeyError("batch needs log_mel, input_features, or features") endpoint = batch.get("endpoint", batch.get("endpoint_bool", batch.get("label"))) if endpoint is None: raise KeyError("batch needs endpoint, endpoint_bool, or label") features = torch.as_tensor(features, dtype=torch.float32, device=device) endpoint = torch.as_tensor(endpoint, dtype=torch.float32, device=device) mask_value = batch.get("attention_mask", batch.get("frame_mask")) if mask_value is None: attention_mask = torch.ones( (features.shape[0], features.shape[-1]), dtype=torch.bool, device=device ) else: attention_mask = torch.as_tensor(mask_value, dtype=torch.bool, device=device) def optional(name: str) -> Tensor | None: value = batch.get(name) return None if value is None else torch.as_tensor(value, dtype=torch.float32, device=device) return features, attention_mask, endpoint, optional("midfiller"), optional("endfiller") class Trainer: def __init__( self, model: nn.Module, config: TrainerConfig | None = None, loss_config: MultiTaskLossConfig | None = None, output_dir: str | Path = "artifacts/run", artifact_metadata: Mapping[str, Any] | None = None, ) -> None: self.model = model self.config = config or TrainerConfig() self.loss = MultiTaskTurnLoss(loss_config) self.output_dir = Path(output_dir) self.artifact_metadata = dict(artifact_metadata or {}) self.device = _resolve_device(self.config.device) self.history: list[dict[str, Any]] = [] def fit(self, train_loader: Iterable, validation_loader: Iterable) -> dict[str, Any]: seed_everything(self.config.seed, self.config.deterministic) self.output_dir.mkdir(parents=True, exist_ok=True) self.model.to(self.device) trainable = [parameter for parameter in self.model.parameters() if parameter.requires_grad] if not trainable: raise ValueError("model has no trainable parameters") optimizer = torch.optim.AdamW( trainable, lr=self.config.learning_rate, weight_decay=self.config.weight_decay, ) scaler = torch.amp.GradScaler( "cuda", enabled=self.config.mixed_precision and self.device.type == "cuda" ) best_score = -float("inf") selected_validation_loss = float("inf") stale_epochs = 0 best_path = self.output_dir / "best.pt" for epoch in range(1, self.config.epochs + 1): train_dataset = getattr(train_loader, "dataset", None) if hasattr(train_dataset, "set_epoch"): train_dataset.set_epoch(epoch - 1) train_metrics = self._train_epoch(train_loader, optimizer, scaler) validation_metrics, labels, probabilities = self.evaluate(validation_loader) operating_point = threshold_at_max_fpr( labels, probabilities, self.config.false_positive_rate_budget, ) record = { "epoch": epoch, "train": train_metrics, "validation": validation_metrics, "operating_point": operating_point, } self.history.append(record) print( json.dumps( { "epoch": epoch, "train_loss": train_metrics["loss"], "validation_loss": validation_metrics["loss"], "validation_roc_auc": validation_metrics["roc_auc"], "validation_average_precision": validation_metrics["average_precision"], "operating_threshold": operating_point["threshold"], "operating_fpr": operating_point["false_positive_rate"], "operating_recall": operating_point["recall"], }, allow_nan=False, ), flush=True, ) validation_loss = float(validation_metrics["loss"]) if self.config.selection_metric == "validation_loss": selection_score = -validation_loss elif self.config.selection_metric == "average_precision": selection_score = float(validation_metrics["average_precision"] or 0.0) elif self.config.selection_metric == "constrained_recall": selection_score = float(operating_point["recall"] or 0.0) else: raise ValueError( "selection_metric must be constrained_recall, average_precision, " "or validation_loss" ) record["selection_metric"] = self.config.selection_metric record["selection_score"] = selection_score improved = selection_score > best_score + self.config.min_improvement if improved: best_score = selection_score selected_validation_loss = validation_loss stale_epochs = 0 self._save_checkpoint( best_path, epoch, optimizer, threshold=float(operating_point["threshold"]), metrics=record, ) else: stale_epochs += 1 self._write_history() if stale_epochs >= self.config.early_stopping_patience: break checkpoint = torch.load(best_path, map_location=self.device, weights_only=False) self.model.load_state_dict(checkpoint["model_state"]) return { "checkpoint": str(best_path), "best_validation_loss": selected_validation_loss, "best_selection_metric": self.config.selection_metric, "best_selection_score": best_score, "threshold": checkpoint["threshold"], "epochs_completed": len(self.history), "history": self.history, } def _train_epoch( self, loader: Iterable, optimizer: torch.optim.Optimizer, scaler: Any ) -> dict[str, float]: self.model.train() optimizer.zero_grad(set_to_none=True) loss_totals = {name: 0.0 for name in ("total", "endpoint", "midfiller", "endfiller")} batches = 0 accumulation = max(1, self.config.gradient_accumulation_steps) for batch_index, batch in enumerate(loader, start=1): features, mask, endpoint, mid, end = _extract_batch(batch, self.device) autocast_enabled = self.config.mixed_precision and self.device.type == "cuda" with torch.autocast(device_type=self.device.type, enabled=autocast_enabled): output: TurnDetectionOutput = self.model(features, mask) losses = self.loss(output, endpoint, mid, end) scaled_loss = losses["total"] / accumulation scaler.scale(scaled_loss).backward() if batch_index % accumulation == 0: scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_( self.model.parameters(), self.config.max_gradient_norm ) scaler.step(optimizer) scaler.update() optimizer.zero_grad(set_to_none=True) for name in loss_totals: loss_totals[name] += float(losses[name].detach().cpu()) batches += 1 if batches == 0: raise ValueError("training loader produced no batches") # Flush a partial accumulation window. if batches % accumulation: scaler.unscale_(optimizer) correction = accumulation / (batches % accumulation) for parameter in self.model.parameters(): if parameter.grad is not None: parameter.grad.mul_(correction) torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.config.max_gradient_norm) scaler.step(optimizer) scaler.update() optimizer.zero_grad(set_to_none=True) return { "loss": loss_totals["total"] / batches, "endpoint_loss": loss_totals["endpoint"] / batches, "midfiller_loss": loss_totals["midfiller"] / batches, "endfiller_loss": loss_totals["endfiller"] / batches, "batches": float(batches), } @torch.no_grad() def evaluate(self, loader: Iterable) -> tuple[dict[str, Any], list[int], list[float]]: self.model.eval() loss_totals = {name: 0.0 for name in ("total", "endpoint", "midfiller", "endfiller")} batches = 0 labels: list[int] = [] probabilities: list[float] = [] for batch in loader: features, mask, endpoint, mid, end = _extract_batch(batch, self.device) output: TurnDetectionOutput = self.model(features, mask) losses = self.loss(output, endpoint, mid, end) for name in loss_totals: loss_totals[name] += float(losses[name].cpu()) batches += 1 labels.extend(int(value) for value in endpoint.detach().cpu().tolist()) probabilities.extend( float(value) for value in torch.sigmoid(output.endpoint_logits).cpu().tolist() ) if batches == 0: raise ValueError("validation loader produced no batches") metrics = binary_classification_metrics(labels, probabilities, threshold=0.5) metrics["loss"] = loss_totals["total"] / batches metrics["endpoint_loss"] = loss_totals["endpoint"] / batches metrics["midfiller_loss"] = loss_totals["midfiller"] / batches metrics["endfiller_loss"] = loss_totals["endfiller"] / batches return metrics, labels, probabilities def _save_checkpoint( self, path: Path, epoch: int, optimizer: torch.optim.Optimizer, threshold: float, metrics: Mapping[str, Any], ) -> None: model_config = ( self.model.model_config() if hasattr(self.model, "model_config") else self.artifact_metadata.get("model_config") ) if not isinstance(model_config, Mapping): raise ValueError("model must expose model_config() for a self-describing checkpoint") payload = { "format_version": 1, "epoch": epoch, "model_config": dict(model_config), "model_state": self.model.state_dict(), "optimizer_state": optimizer.state_dict(), "trainer_config": asdict(self.config), "threshold": threshold, "metrics": dict(metrics), "metadata": self.artifact_metadata, } temporary = path.with_suffix(path.suffix + ".tmp") torch.save(payload, temporary) os.replace(temporary, path) def _write_history(self) -> None: path = self.output_dir / "history.json" temporary = path.with_suffix(".json.tmp") temporary.write_text(json.dumps(self.history, indent=2, allow_nan=False), encoding="utf-8") os.replace(temporary, path)