src/turn_detection/training/trainer.py

"""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)