File size: 1,877 Bytes
b5731a6 | 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 | """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)
|