| """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 |
| from nirnay.model import QTYPES |
| from nirnay.render import render_question, render_state |
|
|
| 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 |
| 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")) |
| 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() |
|
|
