File size: 1,008 Bytes
ff5f59d | 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 | """Shared helpers: tokenizer, bucket choice, scoring API over any backend that maps ids[B,T] -> lp[B,T-1]."""
import json, os
import numpy as np
from tokenizers import Tokenizer
from qdec import HF
TOK = Tokenizer.from_file(f"{HF}/tokenizer.json")
BOS = TOK.token_to_id("<bos>")
PAD = TOK.token_to_id("<pad>")
BUCKETS = [16, 32, 64, 96]
CACHE = os.environ.get("MCBPMF_LM_WORK", ".") + "/results/walk2/dec_cache.v2.f3c75593dbd9.7ab4ee7.fdde1635a116.jsonl"
def ids(text):
return [BOS] + TOK.encode(text).ids
def bucket(n):
for b in BUCKETS:
if n <= b:
return b
raise ValueError(f"sequence of {n} tokens > {BUCKETS[-1]}")
def pack(seqs, B, T):
x = np.full((B, T), PAD, dtype=np.int32)
for i, s in enumerate(seqs):
x[i, :len(s)] = s
return x
def sums(lp, seqs):
return [float(np.asarray(lp[i, :len(s) - 1], dtype=np.float64).sum()) for i, s in enumerate(seqs)]
def load_cache():
return [json.loads(l) for l in open(CACHE, encoding="utf-8")]
|