utils/config.py

from __future__ import annotations

import json
from copy import deepcopy
from pathlib import Path
from typing import Any

import yaml


ROOT = Path(__file__).resolve().parents[1]


def load_config(path: str | Path) -> dict[str, Any]:
    path = Path(path)
    if not path.is_absolute():
        path = ROOT / path
    with path.open("r", encoding="utf-8") as f:
        if path.suffix.lower() == ".json":
            return json.load(f)
        return yaml.safe_load(f) or {}


def deep_merge(base: dict[str, Any], override: dict[str, Any]) -> dict[str, Any]:
    merged = deepcopy(base)
    for key, value in override.items():
        if isinstance(value, dict) and isinstance(merged.get(key), dict):
            merged[key] = deep_merge(merged[key], value)
        else:
            merged[key] = deepcopy(value)
    return merged


def load_dataset_config(name_or_path: str) -> dict[str, Any]:
    path = Path(name_or_path)
    if path.suffix:
        return load_config(path)
    return load_config(Path("configs/datasets") / f"{name_or_path}.yaml")


def load_model_registry(path: str | Path = "configs/models/registry.yaml") -> dict[str, Any]:
    return load_config(path).get("models", {})