src/mini_transformer/model_loader.py

"""Helpers for loading model configurations from local trained_models folders."""

from __future__ import annotations

import os
import shutil
from collections.abc import Sequence
from pathlib import Path

from huggingface_hub import snapshot_download
from hydra import compose, initialize_config_dir
from hydra.core.global_hydra import GlobalHydra
from omegaconf import DictConfig

MODELS_ENV = "MINI_TRANSFORMER_MODELS"
MODEL_NAME_ENV = "MINI_TRANSFORMER_MODEL_NAME"
CONFIG_DIR_ENV = "MINI_TRANSFORMER_CONFIG_DIR"
CONFIG_NAME_ENV = "MINI_TRANSFORMER_CONFIG_NAME"
DEFAULT_CONFIG_NAME = "config_inference"
CONFIG_FILENAME = f"{DEFAULT_CONFIG_NAME}.yaml"


def get_models_root() -> Path:
    """Return the directory where local models are stored."""

    env_root = os.environ.get(MODELS_ENV)
    if env_root:
        return Path(env_root).expanduser().resolve()
    return (Path.cwd() / "trained_models").resolve()


def ensure_models_root() -> Path:
    root = get_models_root()
    root.mkdir(parents=True, exist_ok=True)
    return root


def list_model_names() -> list[str]:
    root = ensure_models_root()
    names: list[str] = []
    for entry in sorted(root.iterdir()):
        if entry.is_dir() and (entry / "configs" / CONFIG_FILENAME).exists():
            names.append(entry.name)
    return names


def _resolve_relative_paths(cfg: DictConfig, base_dir: Path) -> None:
    def _resolve(value: str | None) -> str | None:
        if not value:
            return value
        path = Path(value)
        if path.is_absolute():
            return str(path)
        parts = path.parts
        if len(parts) >= 2 and parts[0] == "trained_models" and parts[1] == base_dir.name:
            path = Path(*parts[2:])
        elif parts and parts[0] == base_dir.name:
            path = Path(*parts[1:])
        return str((base_dir / path).resolve())

    if "model" in cfg:
        for key in ("best_checkpoint_path", "latest_checkpoint_path"):
            if key in cfg.model:
                cfg.model[key] = _resolve(cfg.model[key])
    if "tokenizer" in cfg and "path" in cfg.tokenizer:
        cfg.tokenizer["path"] = _resolve(cfg.tokenizer["path"])
    if "runtime" in cfg:
        for key in ("output_dir", "data_dir", "cache_dir", "tokenizer_dir", "checkpoint_path"):
            if key in cfg.runtime:
                cfg.runtime[key] = _resolve(cfg.runtime[key])


def compose_model_config(
    model_name: str,
    *,
    config_name: str = DEFAULT_CONFIG_NAME,
    overrides: Sequence[str] | None = None,
    job_name: str | None = None,
) -> DictConfig:
    """Compose the Hydra config for a given model folder."""

    models_root = ensure_models_root()
    model_dir = models_root / model_name
    config_dir = model_dir / "configs"
    if not config_dir.is_dir():
        raise FileNotFoundError(
            f"Config directory not found for model '{model_name}' at {config_dir}."
        )

    overrides = list(overrides or [])
    job = job_name or f"mini_transformer_{model_name}"
    if GlobalHydra.instance().is_initialized():
        GlobalHydra.instance().clear()
    with initialize_config_dir(config_dir=str(config_dir), job_name=job, version_base=None):
        cfg = compose(config_name=config_name, overrides=overrides)
    _resolve_relative_paths(cfg, model_dir)
    return cfg


def compose_config_from_dir(
    config_dir: str,
    *,
    config_name: str = DEFAULT_CONFIG_NAME,
    overrides: Sequence[str] | None = None,
    job_name: str = "mini_transformer_cli",
) -> DictConfig:
    overrides = list(overrides or [])
    config_path = Path(config_dir).resolve()
    base_dir = config_path.parent if config_path.name == "configs" else config_path
    if GlobalHydra.instance().is_initialized():
        GlobalHydra.instance().clear()
    with initialize_config_dir(config_dir=str(config_path), job_name=job_name, version_base=None):
        cfg = compose(config_name=config_name, overrides=overrides)
    _resolve_relative_paths(cfg, base_dir)
    return cfg


__all__ = [
    "CONFIG_FILENAME",
    "DEFAULT_CONFIG_NAME",
    "MODELS_ENV",
    "MODEL_NAME_ENV",
    "CONFIG_DIR_ENV",
    "CONFIG_NAME_ENV",
    "compose_model_config",
    "compose_config_from_dir",
    "ensure_models_root",
    "get_models_root",
    "list_model_names",
    "snapshot_to_local",
]


def snapshot_to_local(
    model_id: str,
    *,
    local_name: str | None = None,
    revision: str | None = None,
    repo_type: str = "model",
    token: str | None = None,
    cache_dir: str | None = None,
    force: bool = False,
) -> Path:
    """Download a repository from the Hugging Face Hub into `trained_models/`."""

    root = ensure_models_root()
    name = local_name or model_id.replace("/", "__")
    destination = root / name

    if destination.exists() and any(destination.iterdir()):
        if not force:
            raise FileExistsError(
                f"Destination {destination} already exists. Use --force to overwrite."
            )
        for child in destination.iterdir():
            if child.is_dir():
                shutil.rmtree(child)
            else:
                child.unlink()

    destination.mkdir(parents=True, exist_ok=True)

    snapshot_download(
        repo_id=model_id,
        repo_type=repo_type,
        local_dir=str(destination),
        local_dir_use_symlinks=False,
        revision=revision,
        token=token,
        cache_dir=cache_dir,
    )
    return destination