File size: 9,621 Bytes
00ee773 | 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 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 | """Run the frozen encoder once over every unique training text and store the token states.
Outputs in train/cache/<tag>/:
states.f16 / questions.f16 / options.f16 token states, rows of d (float16 memmaps)
*_mid.f16 the same at layer NIRNAY_SAVE_MID (what deep fusion trains on)
*_off.npy [n+1] row offsets, *_vec.npy [n, d] pooled L2-normalised embeddings (float32)
index.jsonl one line per record: state / question / option ids (canonical option order), gold, type, src, split
meta.json backbone, limits, counts
Options for big backbones:
NIRNAY_STATE_MID_ONLY=1 states stop at the mid layer (deep fusion never needs their frozen top layers)
NIRNAY_SHARD=i/n run n processes (one per CPU socket); each encodes every n-th state (by length)
into the same memory-mapped file; shard 0 also does questions/options and the index
python train\\precompute.py
"""
import json
import os
import sys
import time
import numpy as np
import torch
HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.dirname(HERE)
sys.path.insert(0, ROOT)
from nirnay.encoder import load_hf_encoder # noqa: E402
from nirnay.model import QTYPES # noqa: E402
from nirnay.render import render_question, render_state # noqa: E402
BACKBONE = os.environ.get("NIRNAY_BACKBONE", "intfloat/multilingual-e5-small")
POOL = os.environ.get("NIRNAY_POOL", "mean")
PREFIX = os.environ.get("NIRNAY_PREFIX", "query: ")
PREFIX = "" if PREFIX.strip().lower() == "none" else PREFIX # Windows drops empty env vars
TAG = os.environ.get("NIRNAY_TAG", "e5s")
REC = os.environ.get("NIRNAY_RECORDS", os.path.join(HERE, "data", "records.jsonl"))
OUTD = os.path.join(HERE, "cache", TAG)
LIM = {"state_tokens": int(os.environ.get("NIRNAY_TRAIN_STATE_TOKENS", "320")), "question_tokens": 64, "option_tokens": 32}
MID = int(os.environ.get("NIRNAY_SAVE_MID", "8")) # also keep this layer's output (0 = off)
STATE_MID_ONLY = os.environ.get("NIRNAY_STATE_MID_ONLY") == "1"
SHARD, NSHARD = (int(x) for x in os.environ.get("NIRNAY_SHARD", "0/1").split("/"))
def wait_for(path, what):
t = time.time()
while not os.path.exists(path):
time.sleep(2)
if time.time() - t > 6 * 3600:
raise SystemExit("timed out waiting for " + what)
def main():
from tokenizers import Tokenizer
torch.set_num_threads(int(os.environ.get("NIRNAY_THREADS", "16")))
os.makedirs(OUTD, exist_ok=True)
if NSHARD > 1:
assert MID and STATE_MID_ONLY, "sharding is supported for NIRNAY_STATE_MID_ONLY=1 runs"
enc = load_hf_encoder(BACKBONE, POOL, PREFIX)
tok = Tokenizer.from_file(os.path.join(enc.source_dir, "tokenizer.json"))
tok.no_truncation()
tok.no_padding()
probe = tok.encode("a", add_special_tokens=True).ids
cls_id, sep_id = probe[0], probe[-1]
pad_id = enc.cfg.pad_id
prefix_ids = tok.encode(PREFIX, add_special_tokens=False).ids if PREFIX else []
print("backbone", BACKBONE, "cls/sep/pad", cls_id, sep_id, pad_id, "prefix", prefix_ids,
"shard %d/%d" % (SHARD, NSHARD), flush=True)
def ids_for(text, budget):
room = budget - 2 - len(prefix_ids)
return [cls_id] + prefix_ids + tok.encode(text, add_special_tokens=False).ids[:max(1, room)] + [sep_id]
states, questions, options = {}, {}, {}
index = []
t0 = time.time()
with open(REC, encoding="utf-8") as f:
for line in f:
r = json.loads(line)
try:
nq, qtext, otexts = render_question("q", r["q"])
except Exception:
continue
perm = sorted(range(len(otexts)), key=lambda i: (otexts[i], i))
st = render_state(r["state"])
si = states.setdefault(st, len(states))
qi = questions.setdefault(qtext, len(questions))
oi = [options.setdefault(otexts[i], len(options)) for i in perm]
index.append({"s": si, "q": qi, "o": oi, "y": perm.index(int(r["y"])), "t": QTYPES[nq["type"]],
"src": r["src"], "split": r["split"]})
print("records %d | unique states %d questions %d options %d | %.0fs" %
(len(index), len(states), len(questions), len(options), time.time() - t0), flush=True)
def encode_all(texts, budget, name, mid_only=False):
seqs = [ids_for(t, budget) for t in texts]
lens = np.array([len(s) for s in seqs], dtype=np.int64)
off = np.zeros(len(seqs) + 1, dtype=np.int64)
off[1:] = np.cumsum(lens)
d = enc.cfg.hidden
rows = (int(off[-1]), d)
sharded = mid_only and NSHARD > 1
created = os.path.join(OUTD, name + "_mid.created")
if sharded and SHARD != 0:
wait_for(created, name + " memmap")
mm = None
mm_mid = np.lib.format.open_memmap(os.path.join(OUTD, name + "_mid.f16.npy"), mode="r+")
assert mm_mid.shape == rows, (mm_mid.shape, rows)
else:
if sharded and os.path.exists(created):
os.remove(created)
mm = None if mid_only else np.lib.format.open_memmap(os.path.join(OUTD, name + ".f16.npy"), mode="w+",
dtype=np.float16, shape=rows)
mm_mid = np.lib.format.open_memmap(os.path.join(OUTD, name + "_mid.f16.npy"), mode="w+", dtype=np.float16,
shape=rows) if MID else None
if sharded:
mm_mid.flush()
open(created, "w").close()
vec = np.zeros((len(seqs), d), dtype=np.float32)
order = np.argsort(lens, kind="stable")
if sharded:
order = order[SHARD::NSHARD]
total = int(lens[order].sum())
i, done, t1, last = 0, 0, time.time(), 0.0
budget_tok = 12288
with torch.no_grad():
while i < len(order):
j = i
while j < len(order) and lens[order[j]] * (j - i + 1) <= max(budget_tok, lens[order[i]]):
j += 1
j = max(j, i + 1)
idx = order[i:j]
Lm = int(lens[idx].max())
ids = torch.full((len(idx), Lm), pad_id, dtype=torch.long)
mask = torch.zeros((len(idx), Lm), dtype=torch.bool)
for rr, t in enumerate(idx):
ids[rr, :lens[t]] = torch.tensor(seqs[t])
mask[rr, :lens[t]] = True
if mid_only:
Mn = enc(ids, mask, stop_layer=MID).numpy().astype(np.float16)
for rr, t in enumerate(idx):
mm_mid[off[t]:off[t + 1]] = Mn[rr, :lens[t]]
else:
if MID:
H, Mh = enc(ids, mask, keep_layer=MID)
Mn = Mh.numpy().astype(np.float16)
else:
H = enc(ids, mask)
V = enc.pool(H, mask)
Hn = H.numpy().astype(np.float16)
for rr, t in enumerate(idx):
mm[off[t]:off[t + 1]] = Hn[rr, :lens[t]]
if MID:
mm_mid[off[t]:off[t + 1]] = Mn[rr, :lens[t]]
vec[idx] = V.numpy()
done += int(lens[idx].sum())
i = j
if time.time() - last > 5 or i >= len(order):
last = time.time()
el = time.time() - t1
print("\r [%s] %d/%d texts %.1fk tok/s ETA %.0f min " %
(name, i, len(order), done / max(el, 1e-9) / 1000,
(total - done) / max(done / max(el, 1e-9), 1e-9) / 60), end="", flush=True)
if mm is not None:
mm.flush()
del mm
if mm_mid is not None:
mm_mid.flush()
del mm_mid
if sharded:
open(os.path.join(OUTD, "%s.done.%d" % (name, SHARD)), "w").close()
if SHARD != 0:
print("\n [%s] shard %d: %d texts, %d tokens in %.1f min" % (name, SHARD, len(order), total, (time.time() - t1) / 60), flush=True)
return
for s in range(NSHARD):
wait_for(os.path.join(OUTD, "%s.done.%d" % (name, s)), "%s shard %d" % (name, s))
np.save(os.path.join(OUTD, name + "_off.npy"), off)
np.save(os.path.join(OUTD, name + "_vec.npy"), vec)
print("\n [%s] %d texts, %d tokens in %.1f min" % (name, len(seqs), int(off[-1]), (time.time() - t1) / 60), flush=True)
if SHARD == 0:
encode_all(list(questions), LIM["question_tokens"], "questions")
encode_all(list(options), LIM["option_tokens"], "options")
encode_all(list(states), LIM["state_tokens"], "states", mid_only=STATE_MID_ONLY)
if SHARD != 0:
return
with open(os.path.join(OUTD, "index.jsonl"), "w", encoding="utf-8") as f:
for r in index:
f.write(json.dumps(r) + "\n")
json.dump({"backbone": BACKBONE, "pool": POOL, "prefix": PREFIX, "limits": LIM, "d": enc.cfg.hidden, "mid_layer": MID,
"state_mid_only": STATE_MID_ONLY,
"n_records": len(index), "n_states": len(states), "n_questions": len(questions),
"n_options": len(options), "special_ids": {"cls": cls_id, "sep": sep_id, "pad": pad_id},
"encoder": enc.cfg.to_dict(), "source_dir": enc.source_dir},
open(os.path.join(OUTD, "meta.json"), "w"), indent=2)
print("done ->", OUTD)
if __name__ == "__main__":
main()
|