95 lines
3.1 KiB
Python
95 lines
3.1 KiB
Python
"""Runtime environment: offline HuggingFace settings, config, device, host.
|
|
|
|
Import this module before anything that imports huggingface_hub, transformers,
|
|
or sentence_transformers: huggingface_hub reads HF_HUB_OFFLINE at import time.
|
|
(Convention copied from ssearch/query_hybrid.py.)
|
|
|
|
Models are downloaded once with network access (`uv run bingham download-models`,
|
|
which sets ART_ONLINE=1); every other command runs offline.
|
|
|
|
Personal files live in the user directory (art/user/, or ART_USER_DIR), kept
|
|
out of the code repository and usually a small private git repository of its
|
|
own. user/config.yaml is merged over config.yaml, and data files are looked up
|
|
in user/ before data/ (see data_file).
|
|
"""
|
|
|
|
import os
|
|
import platform
|
|
import socket
|
|
from functools import cache
|
|
from pathlib import Path
|
|
|
|
import yaml
|
|
|
|
from art._paths import ROOT
|
|
|
|
|
|
USER_DIR = Path(os.environ.get("ART_USER_DIR", ROOT / "user"))
|
|
|
|
|
|
def _merge(base: dict, over: dict) -> dict:
|
|
"""Recursive dict merge; values in `over` win, lists are replaced, not appended."""
|
|
out = dict(base)
|
|
for key, value in over.items():
|
|
out[key] = _merge(out[key], value) if isinstance(value, dict) and isinstance(out.get(key), dict) else value
|
|
return out
|
|
|
|
|
|
@cache
|
|
def config() -> dict:
|
|
"""config.yaml, with user/config.yaml merged over it if present."""
|
|
with open(ROOT / "config.yaml") as f:
|
|
cfg = yaml.safe_load(f)
|
|
user_cfg = USER_DIR / "config.yaml"
|
|
if user_cfg.exists():
|
|
cfg = _merge(cfg, yaml.safe_load(user_cfg.read_text()) or {})
|
|
return cfg
|
|
|
|
|
|
_warned: set[str] = set()
|
|
|
|
|
|
def data_file(name: str) -> Path:
|
|
"""user/<name>, else data/<name>, else data/examples/<name> (with a one-time notice)."""
|
|
for candidate in (USER_DIR / name, ROOT / "data" / name):
|
|
if candidate.exists():
|
|
return candidate
|
|
example = ROOT / "data" / "examples" / name
|
|
if example.exists():
|
|
if name not in _warned:
|
|
_warned.add(name)
|
|
print(f" note: no {USER_DIR.name}/{name}; using the example in data/examples/ "
|
|
f"(see `bingham init`)")
|
|
return example
|
|
raise FileNotFoundError(f"{name} not found in {USER_DIR}, data/, or data/examples/")
|
|
|
|
|
|
def contact_email() -> str | None:
|
|
"""Contact address sent to arXiv and Crossref (Crossref's "polite pool"); set in user/config.yaml."""
|
|
return (config().get("contact") or {}).get("email") or None
|
|
|
|
|
|
def path(key: str):
|
|
"""A path from config.yaml `paths`, resolved against ROOT."""
|
|
return ROOT / config()["paths"][key]
|
|
|
|
|
|
os.environ["TOKENIZERS_PARALLELISM"] = "false"
|
|
os.environ["TRANSFORMERS_VERBOSITY"] = "error" # token counts over 512 are reported by embed.py
|
|
os.environ["HF_HUB_CACHE"] = str(path("models"))
|
|
os.environ["HF_HUB_OFFLINE"] = "0" if os.environ.get("ART_ONLINE") == "1" else "1"
|
|
|
|
|
|
@cache
|
|
def device() -> str:
|
|
"""cuda, then mps, then cpu."""
|
|
import torch
|
|
if torch.cuda.is_available():
|
|
return "cuda"
|
|
if torch.backends.mps.is_available():
|
|
return "mps"
|
|
return "cpu"
|
|
|
|
|
|
def host() -> str:
|
|
return f"{socket.gethostname()} ({platform.system()} {platform.machine()})"
|