seeker
40 lines
933 B
Python
40 lines
933 B
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import random
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import torch
|
|
import yaml
|
|
|
|
|
|
def load_config(path: str | Path) -> dict:
|
|
with Path(path).open(encoding="utf-8") as stream:
|
|
return yaml.safe_load(stream)
|
|
|
|
|
|
def seed_everything(seed: int) -> None:
|
|
random.seed(seed)
|
|
np.random.seed(seed)
|
|
torch.manual_seed(seed)
|
|
|
|
|
|
def choose_device(name: str = "auto") -> torch.device:
|
|
if name != "auto":
|
|
return torch.device(name)
|
|
if torch.cuda.is_available():
|
|
return torch.device("cuda")
|
|
if torch.backends.mps.is_available():
|
|
return torch.device("mps")
|
|
return torch.device("cpu")
|
|
|
|
|
|
def write_json(path: str | Path, value: object) -> None:
|
|
output = Path(path)
|
|
output.parent.mkdir(parents=True, exist_ok=True)
|
|
output.write_text(
|
|
json.dumps(value, indent=2, ensure_ascii=False) + "\n",
|
|
encoding="utf-8",
|
|
)
|