File size: 1,532 Bytes
35d483e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 | """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
|