File size: 1,623 Bytes
6dd9839 | 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 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 | #!/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")
|