| """tokenise blocks, codec codes (cached), pair/prefix arrays per variant.""" |
| import os, json, numpy as np, mlx.core as mx |
| from gen2 import * |
| def tok_block(t): return (bpe.encode(t)[:L - 1] + [EOS]) |
| def load_split(name, fc, cache=True, limit=0): |
| S = json.load(open(f"{DATA}/{name}.json")) |
| if limit: S, cache = S[:limit], False |
| ids, sidx, bidx = [], [], [] |
| for si, s in enumerate(S): |
| for bi, b in enumerate(s["blocks"]): ids.append(tok_block(b)); sidx.append(si); bidx.append(bi) |
| arr = np.array([i + [PAD] * (L - len(i)) for i in ids], np.int32) |
| f = f"{DATA}/codes_{name}.npy" |
| if cache and os.path.exists(f) and len(np.load(f, mmap_mode='r')) == len(arr): codes = np.load(f) |
| else: |
| out = [] |
| for i in range(0, len(arr), 128): |
| cs = fc.encode(mx.array(arr[i:i + 128])); out.append(np.array(mx.concatenate(cs, 1)).astype(np.int16)); mx.clear_cache() |
| codes = np.concatenate(out) |
| if cache: np.save(f, codes) |
| return S, arr, codes, np.array(sidx), np.array(bidx) |
| def pairs(v, S, arr, codes, sidx, bidx, ntable_from="gold"): |
| """one pair per (story, block n>=1): prefix array (P,), target block row.""" |
| start = {} |
| for r, (s, b) in enumerate(zip(sidx, bidx)): |
| if b == 0: start[s] = r |
| import re |
| SR = re.compile(r"(?<![\w'])(?:John|Mary|David|Sarah|James|Peter|Emma|Mark)(?![\w'])") |
| P, T = [], [] |
| for s, st in enumerate(S): |
| seen, r0 = [], start[s] |
| for n in range(1, len(st["blocks"])): |
| for m in SR.findall(st["blocks"][n - 1]): |
| if m not in seen: seen.append(m) |
| prev = [t for t in arr[r0 + n - 1].tolist() if t not in (EOS, PAD)] |
| hist = [codes[r0 + j][:CPB] for j in range(max(0, n - MB), n)] |
| P.append(make_prefix(v, prev, seen, hist)); T.append(r0 + n) |
| return np.stack(P), np.array(T) |
|
|
