| """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")] | |