48 lines
1.6 KiB
Python
48 lines
1.6 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import argparse
|
||
|
|
import os
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
|
||
|
|
DEFAULT_KEY_PATH = Path("KEY.md")
|
||
|
|
|
||
|
|
|
||
|
|
def load_key_file(path: Path) -> dict[str, str]:
|
||
|
|
if not path.exists():
|
||
|
|
return {}
|
||
|
|
loaded: dict[str, str] = {}
|
||
|
|
for raw_line in path.read_text(encoding="utf-8").splitlines():
|
||
|
|
line = raw_line.strip()
|
||
|
|
if not line or line.startswith("#") or "=" not in line:
|
||
|
|
continue
|
||
|
|
key, value = line.split("=", 1)
|
||
|
|
loaded[key.strip()] = value.strip()
|
||
|
|
return loaded
|
||
|
|
|
||
|
|
|
||
|
|
def ensure_tokens(args: argparse.Namespace) -> None:
|
||
|
|
primary_key_path = Path(getattr(args, "key_path", DEFAULT_KEY_PATH))
|
||
|
|
if not primary_key_path.exists():
|
||
|
|
parent_key = primary_key_path.parent.parent / primary_key_path.name
|
||
|
|
if parent_key.exists():
|
||
|
|
primary_key_path = parent_key
|
||
|
|
values = load_key_file(primary_key_path)
|
||
|
|
|
||
|
|
if not getattr(args, "hf_token", None):
|
||
|
|
args.hf_token = os.getenv("HF_TOKEN") or values.get("HF_TOKEN")
|
||
|
|
|
||
|
|
modelhub_token = getattr(args, "modelhub_token", None)
|
||
|
|
if not modelhub_token:
|
||
|
|
modelhub_token = os.getenv("MODELHUB_XC_TOKEN") or os.getenv("XC_TOKEN") or values.get("MODELHUB_XC_TOKEN") or values.get("XC_TOKEN")
|
||
|
|
|
||
|
|
args.modelhub_token = modelhub_token
|
||
|
|
|
||
|
|
if args.hf_token:
|
||
|
|
os.environ["HF_TOKEN"] = args.hf_token
|
||
|
|
if args.modelhub_token:
|
||
|
|
os.environ["MODELHUB_XC_TOKEN"] = args.modelhub_token
|
||
|
|
os.environ["XC_TOKEN"] = args.modelhub_token
|
||
|
|
if not args.modelhub_token and not os.getenv("STRATEGY_ID"):
|
||
|
|
raise ValueError("XC_TOKEN is required, either via environment or KEY.md")
|