Source code for cann.pretrained

"""Load pretrained surrogate artifacts bundled with the package."""

from __future__ import annotations

import glob
import os

import torch

from cann.config.model_config import ModelConfig
from cann.engine.training import SurrogateModel, TrainedModel
from cann.surrogate.mlp import build_mlp_network


def _default_artifact_dir(subdir: str) -> str:
    return os.path.join(os.path.dirname(os.path.abspath(__file__)), "artifacts", subdir)


def _load_pretrained(
    *,
    artifact_dir: str | None,
    default_subdir: str,
    config_filename: str,
) -> TrainedModel:
    if artifact_dir is None:
        artifact_dir = _default_artifact_dir(default_subdir)

    if not os.path.exists(artifact_dir):
        raise FileNotFoundError(f"Artifact directory not found: {artifact_dir}")

    config_path = os.path.join(artifact_dir, config_filename)
    if not os.path.exists(config_path):
        raise FileNotFoundError(
            f"Configuration file '{config_filename}' not found in {artifact_dir}"
        )

    model_config = ModelConfig.from_yaml(config_path)

    pattern = os.path.join(artifact_dir, "model_state_*.pth")
    files = glob.glob(pattern)
    if not files:
        raise FileNotFoundError(f"No model state files found matching {pattern}")

    latest_file = max(files, key=os.path.getmtime)

    device = "cuda" if torch.cuda.is_available() else "cpu"
    state_dict = torch.load(latest_file, map_location=device)

    if model_config.train_defaults is None:
        raise ValueError(
            "Model config must have train_defaults to rebuild model architecture."
        )

    network = build_mlp_network(model_config, model_config.train_defaults)
    network.load_state_dict(state_dict)
    network.to(device)

    surrogate = SurrogateModel(network=network, model_cfg=model_config)
    surrogate.eval()

    return TrainedModel(
        surrogate=surrogate,
        training_config=model_config.train_defaults,
        model_config=model_config,
        training_metrics={},
        path=latest_file,
    )




def PretrainedHeston(artifact_dir: str | None = None) -> TrainedModel:
    """Load the bundled pretrained Heston implied-volatility surrogate.

    Args:
        artifact_dir (str | None): Optional path to an artifact directory. If not provided,
            uses the package-bundled ``cann/artifacts/heston_iv_v1``.

    Returns:
        TrainedModel: Loaded trained model wrapper.

    Raises:
        FileNotFoundError: If the artifact directory/config/state dict cannot be found.
        ValueError: If the model config does not include ``train_defaults``.
    """

    return _load_pretrained(
        artifact_dir=artifact_dir,
        default_subdir="heston_iv_v1",
        config_filename="heston_iv.yaml",
    )





def PretrainedRoughHeston(artifact_dir: str | None = None) -> TrainedModel:
    """Load the bundled pretrained Rough Heston surrogate.

    Args:
        artifact_dir (str | None): Optional path to an artifact directory. If not provided,
            uses the package-bundled ``cann/artifacts/rough_heston_v1``.

    Returns:
        TrainedModel: Loaded trained model wrapper.

    Raises:
        FileNotFoundError: If the artifact directory/config/state dict cannot be found.
        ValueError: If the model config does not include ``train_defaults``.
    """

    return _load_pretrained(
        artifact_dir=artifact_dir,
        default_subdir="rough_heston_v1",
        config_filename="rough_heston.yaml",
    )



__all__ = [
    "PretrainedHeston",
    "PretrainedRoughHeston",
]