File size: 4,023 Bytes
f930dac | 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 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 | """Shared constants, paths and IO utilities.
Every path is environment-overridable, and every default is relative to this
repository, so nothing depends on the machine the study was run on.
PNS_DATA checkout of the companion dataset repo (default: ./data)
PNS_CKPT published checkpoints (default: ./checkpoints)
PNS_EVAL where evaluation output is written (default: ./results/eval)
PNS_TOKENIZER tokenizer.json (default: $PNS_DATA/benchmark/tokenizer.json)
`PNS_DATA` must point at the root of the `pns-world` dataset repo, i.e. the
directory containing `data/`, `benchmark/` and `manifests/`.
GPU selection is left entirely to `CUDA_VISIBLE_DEVICES`; nothing here binds a
device.
"""
from __future__ import annotations
import hashlib
import json
import os
import tempfile
from pathlib import Path
# ---------------------------------------------------------------- paths
PNS_ROOT = Path(os.environ.get("PNS_ROOT", Path(__file__).resolve().parents[2]))
def data_root() -> Path:
"""Root of the pns-world dataset checkout (the dir holding data/ and
benchmark/). Set PNS_DATA to point at it; otherwise it is located by
looking in the usual places relative to this repository."""
env = os.environ.get("PNS_DATA")
if env:
return Path(env)
for cand in (PNS_ROOT / "data", PNS_ROOT, PNS_ROOT.parent / "pns-world"):
if (cand / "data" / "e1_original").is_dir() or (cand / "manifests").is_dir():
return cand
return PNS_ROOT / "data"
def shards_root() -> Path:
"""Directory holding the split directories (see pns.data.view.SPLIT_DIRS)."""
return data_root() / "data"
def eval_root() -> Path:
return Path(os.environ.get("PNS_EVAL", PNS_ROOT / "results" / "eval"))
def ckpt_root() -> Path:
return Path(os.environ.get("PNS_CKPT", PNS_ROOT / "checkpoints"))
def logs_root() -> Path:
return Path(os.environ.get("PNS_LOGS", PNS_ROOT / "logs"))
def tokenizer_path() -> Path:
return Path(os.environ.get("PNS_TOKENIZER",
data_root() / "benchmark" / "tokenizer.json"))
# ---------------------------------------------------------------- atomic io
def atomic_write_json(path: Path | str, obj) -> None:
path = Path(path)
path.parent.mkdir(parents=True, exist_ok=True)
fd, tmp = tempfile.mkstemp(dir=path.parent, prefix=path.name + ".tmp")
try:
with os.fdopen(fd, "w") as f:
json.dump(obj, f, indent=1, sort_keys=True, default=str)
f.flush()
os.fsync(f.fileno())
os.replace(tmp, path)
dirfd = os.open(path.parent, os.O_RDONLY)
try:
os.fsync(dirfd)
finally:
os.close(dirfd)
except BaseException:
if os.path.exists(tmp):
os.unlink(tmp)
raise
def sha256_file(path: Path | str) -> str:
h = hashlib.sha256()
with open(path, "rb") as f:
for chunk in iter(lambda: f.read(1 << 20), b""):
h.update(chunk)
return h.hexdigest()
def sha256_text(text: str) -> str:
return hashlib.sha256(text.encode()).hexdigest()
# ---------------------------------------------------- metadata firewall classes
# Keys a model forward() may consume vs keys that exist only for supervision or
# for eval bucketing. tests assert disjointness and that the model consumes
# only MODEL_INPUT_KEYS.
MODEL_INPUT_KEYS = frozenset({
"tokens", "ev_bounds", "etype", "dt",
"rec_store", "rec_kind", "rec_ent", "rec_key", "rec_val_toks", "rec_key_toks",
"rec_birth_ev", "cache_live", "cache_adds",
"enum_legal", # the legal answer set is part of the question, runtime-visible
})
SUPERVISION_KEYS = frozenset({
"mode_gold", "family", "enum_gold", "ptr_gold_slot", "op_gold", "op_arg_slots",
"gold_val_hash",
})
EVAL_METADATA_KEYS = frozenset({
"delay", "kappa", "gold_age_rank", "chain_len", "chain_pos", "corrected",
"world_seed", "lifetime_id", "q_clock",
})
SCHEMA_VERSION = "PNS_V1"
|