api/common.py

#!/usr/bin/env python3
"""Shared path helpers and env loading."""
from __future__ import annotations

import os
from pathlib import Path


PKG_ROOT = Path(__file__).resolve().parent.parent


def load_env(env_path: Path | None = None) -> None:
    path = env_path or (PKG_ROOT / "configs" / "paths.env")
    if not path.is_file():
        example = PKG_ROOT / "configs" / "paths.env.example"
        raise FileNotFoundError(
            f"Missing {path}. Copy {example} to paths.env and set VESM paths."
        )
    for raw in path.read_text(encoding="utf-8").splitlines():
        line = raw.strip()
        if not line or line.startswith("#") or "=" not in line:
            continue
        key, val = line.split("=", 1)
        key, val = key.strip(), val.strip().strip('"').strip("'")
        if key and key not in os.environ:
            os.environ[key] = val


def resolve(rel_or_abs: str) -> Path:
    p = Path(rel_or_abs)
    if p.is_absolute():
        return p
    return (PKG_ROOT / p).resolve()


def vesm_paths() -> tuple[Path, Path]:
    base = os.environ.get("VESM_BASE_MODEL_DIR")
    weights = os.environ.get("VESM_WEIGHTS")
    root = os.environ.get("VESM_ROOT")
    if base and weights:
        return Path(base), Path(weights)
    if not root:
        raise RuntimeError(
            "Set VESM_ROOT or VESM_BASE_MODEL_DIR+VESM_WEIGHTS in configs/paths.env"
        )
    root_p = Path(root)
    return (
        root_p / "models" / "base" / "facebook_esm2_t36_3B_UR50D",
        root_p / "models" / "weights" / "VESM_3B.pth",
    )


def device() -> str:
    return os.environ.get("DEVICE", "cuda:0")