src/turn_detection/models/factory.py

"""Model construction and checkpoint loading."""

from __future__ import annotations

from collections.abc import Mapping
from pathlib import Path
from typing import Any

import torch
from torch import nn

from .tiny_tcn import TinyTCNConfig, TinyTurnDetector


def build_model(config: Mapping[str, Any]) -> nn.Module:
    model_type = str(config.get("type", "tiny_tcn"))
    if model_type in {"tiny_tcn", "student", "tiny"}:
        return TinyTurnDetector(TinyTCNConfig.from_mapping(config))
    if model_type in {"whisper_teacher", "whisper", "teacher"}:
        from .whisper_teacher import WhisperTeacherConfig, WhisperTurnTeacher

        return WhisperTurnTeacher(WhisperTeacherConfig.from_mapping(config))
    raise ValueError(f"unknown model type: {model_type!r}")


def load_model_checkpoint(
    checkpoint_path: str | Path, map_location: str | torch.device = "cpu"
) -> tuple[nn.Module, Mapping[str, Any]]:
    """Load a self-describing training checkpoint."""

    checkpoint = torch.load(checkpoint_path, map_location=map_location, weights_only=False)
    if not isinstance(checkpoint, Mapping) or "model_state" not in checkpoint:
        raise ValueError("checkpoint must contain model_state and model_config")
    model_config = checkpoint.get("model_config")
    if not isinstance(model_config, Mapping):
        raise ValueError("checkpoint does not contain a valid model_config")
    model = build_model(model_config)
    model.load_state_dict(checkpoint["model_state"], strict=True)
    return model, checkpoint