nx_nofloat_tok.nx source
↩ module page · 182 lines · 9015 B
1// nx_nofloat_tok.nx -- canonical sovereign BPE tokenizer library for the GGUF vocab (ONE source of truth).
2// text <-> token ids using tokenizer.ggml.tokens (vocab) + tokenizer.ggml.merges (ranked merge rules).
3// tk_find_id(target) -> id (string -> id; search vocab)
4// tk_decode_off(id) -> byte offset of the id's vocab string (read len/ptr via nx_gguf_meta)
5// tk_bpe_encode(input) -> tokens + ids (single-char start; merge lowest-rank adjacent pair until none)
6// Gates import this (DRY -- no copy-paste tokenizer code). ASCII inputs only for now (byte-level pretokenize
7// is identity for printable ASCII; the GPT-2 byte->unicode map for spaces/non-ASCII is the follow-on).
8// No main (pure library). license_tier: ORIGINAL
9import "nx_syscalls.nx"
10import "nx_tier.nx"
11import "nx_le.nx"
12import "nx_tensor.nx"
13import "nx_gguf.nx"
14import "nx_gguf_meta.nx"
15
16func tk_bytes_eq(a: *u8, b: *u8, n: i64) -> i64 { var i: i64=0; while i<n { if a[i]!=b[i] { return 0 } i=i+1 } return 1 }
17
18// token-string -> vocab id (search the length-prefixed tokens array), or -1.
19func tk_find_id(buf: *u8, first: i64, vocab: i64, target: *u8, tlen: i64) -> i64 {
20 var off: i64=first; var i: i64=0
21 while i<vocab { let sl: i64=nx_gguf_meta_read_string_len(buf, off); if sl==tlen { let sp: *u8=nx_gguf_meta_read_string_ptr(buf, off); if tk_bytes_eq(sp, target, tlen)==1 { return i } } off=off+8+sl; i=i+1 }
22 return 0 - 1
23}
24// byte offset of vocab element `tid` (caller reads string len/ptr there).
25func tk_decode_off(buf: *u8, first: i64, tid: i64) -> i64 {
26 var off: i64=first; var i: i64=0
27 while i<tid { let sl: i64=nx_gguf_meta_read_string_len(buf, off); off=off+8+sl; i=i+1 }
28 return off
29}
30// rank of merge rule "X Y" in the merges array, or -1.
31func tk_merge_rank(buf: *u8, mfirst: i64, nm: i64, xp: *u8, xl: i64, yp: *u8, yl: i64, tmp: *u8) -> i64 {
32 var p: i64=0; var i: i64=0
33 while i<xl { tmp[p]=xp[i]; p=p+1; i=i+1 }
34 tmp[p]=32 as u8; p=p+1
35 i=0; while i<yl { tmp[p]=yp[i]; p=p+1; i=i+1 }
36 let tl: i64=p
37 var off: i64=mfirst; var j: i64=0
38 while j<nm { let ml: i64=nx_gguf_meta_read_string_len(buf, off); if ml==tl { let mp: *u8=nx_gguf_meta_read_string_ptr(buf, off); if tk_bytes_eq(mp, tmp, tl)==1 { return j } } off=off+8+ml; j=j+1 }
39 return 0 - 1
40}
41// byte b -> its GPT-2 byte-level char, UTF-8-encoded into out; returns length (1 or 2).
42// printable bytes (33-126, 161-172, 174-255) map to themselves; the rest map to U+0100+rank (e.g. space 0x20 -> U+0120 'Ġ').
43func tk_byte_char(b: i64, out: *u8) -> i64 {
44 var printable: i64=0
45 if b>=33 { if b<=126 { printable=1 } }
46 if b>=161 { if b<=172 { printable=1 } }
47 if b>=174 { if b<=255 { printable=1 } }
48 var cp: i64=0
49 if printable==1 { cp=b } else { var rank: i64=0; if b<=32 { rank=b } else { if b<=160 { rank=33+(b-127) } else { rank=67 } } cp=256+rank }
50 if cp<128 { out[0]=cp as u8; return 1 }
51 out[0]=(192|(cp>>6)) as u8; out[1]=(128|(cp&63)) as u8
52 return 2
53}
54// ---- fast lookup tables (2026-07-15): FNV-1a hash over the length-prefixed meta arrays, built per
55// encode call (two O(count) walks, ~ms), probed O(1). Replaces the per-pair LINEAR scans of the 151k
56// merges + per-token vocab walks that made KB-scale prompts take ~an hour (measured: 4KB mini-pack
57// prompt stuck >40min single-core in tokenize; 805B was ~150s). Byte-identical semantics: exact
58// byte-equality on probe, duplicate entries resolve to the LOWEST index (insertion order = probe
59// order under linear probing), merge choice stays "lowest rank, first position" -- proven by the
60// tokenizer KAT gate.
61const TK_TCAP: i64 = 262144
62
63func tk_h64(p: *u8, n: i64) -> i64 {
64 var h: i64 = 1469598103934665603
65 var i: i64 = 0
66 while i < n {
67 let b: i64 = p[i] as i64
68 h = h ^ b
69 h = h * 1099511628211
70 i = i + 1
71 }
72 return h
73}
74
75// insert all count elements of a length-prefixed string array into tab (cap TK_TCAP slots x [off,val])
76func tk_tab_build(buf: *u8, first: i64, count: i64, tab: *i64) -> i64 {
77 var off: i64 = first
78 var j: i64 = 0
79 while j < count {
80 let sl: i64 = nx_gguf_meta_read_string_len(buf, off)
81 let sp: *u8 = nx_gguf_meta_read_string_ptr(buf, off)
82 var idx: i64 = tk_h64(sp, sl) & (TK_TCAP - 1)
83 var placing: i64 = 1
84 while placing == 1 {
85 let i2: i64 = idx * 2
86 if tab[i2] == 0 { tab[i2] = off; tab[i2 + 1] = j; placing = 0 } else { idx = (idx + 1) & (TK_TCAP - 1) }
87 }
88 off = off + 8 + sl
89 j = j + 1
90 }
91 return 0
92}
93
94func tk_tab_get(buf: *u8, tab: *i64, q: *u8, ql: i64) -> i64 {
95 var idx: i64 = tk_h64(q, ql) & (TK_TCAP - 1)
96 var probes: i64 = 0
97 while probes < TK_TCAP {
98 let i2: i64 = idx * 2
99 let off: i64 = tab[i2]
100 if off == 0 { return 0 - 1 }
101 let sl: i64 = nx_gguf_meta_read_string_len(buf, off)
102 if sl == ql {
103 let sp: *u8 = nx_gguf_meta_read_string_ptr(buf, off)
104 if tk_bytes_eq(sp, q, ql) == 1 { return tab[i2 + 1] }
105 }
106 idx = (idx + 1) & (TK_TCAP - 1)
107 probes = probes + 1
108 }
109 return 0 - 1
110}
111
112// rank of the pair "X Y" via the merges table; tmp holds the joined bytes (>500B pairs cannot match
113// any real merge rule -> -1 without probing, same as the linear scan's no-match)
114func tk_rank_fast(buf: *u8, mtab: *i64, xp: *u8, xl: i64, yp: *u8, yl: i64, tmp: *u8) -> i64 {
115 if xl + yl + 1 > 500 { return 0 - 1 }
116 var p: i64 = 0
117 var i: i64 = 0
118 while i < xl { tmp[p] = xp[i]; p = p + 1; i = i + 1 }
119 tmp[p] = 32 as u8
120 p = p + 1
121 i = 0
122 while i < yl { tmp[p] = yp[i]; p = p + 1; i = i + 1 }
123 return tk_tab_get(buf, mtab, tmp, p)
124}
125
126// BPE-encode input[0..ilen) -> tokens (tok_ptr/tok_len) + ids; returns ntok. Byte-level pretokenize first.
127// O(n log-ish): hash tables for rank/id lookups + an incremental pair-rank cache (a merge only changes
128// the two pairs adjacent to it; everything else keeps its cached rank).
129func tk_bpe_encode(buf: *u8, mfirst: i64, nm: i64, vfirst: i64, vocab: i64, input: *u8, ilen: i64, tok_ptr: *i64, tok_len: *i64, ids: *i64) -> i64 {
130 if nm >= 200000 { return 0 - 1 }
131 if vocab >= 200000 { return 0 - 1 }
132 let tmp: *u8 = sys_mmap(512)
133 // byte-level pretokenize: each input byte -> its byte-level char (1-2 UTF-8 bytes) in a pool
134 let pool: *u8 = sys_mmap(ilen*4 + 16)
135 var ntok: i64=0; var pp: i64=0; var i: i64=0
136 while i<ilen { let cl: i64=tk_byte_char(input[i] as i64, ((pool as i64)+pp) as *u8); tok_ptr[ntok]=(pool as i64)+pp; tok_len[ntok]=cl; pp=pp+cl; ntok=ntok+1; i=i+1 }
137 let ntok0: i64 = ntok
138 let mtab: *i64 = sys_mmap(TK_TCAP * 16) as *i64
139 let vtab: *i64 = sys_mmap(TK_TCAP * 16) as *i64
140 tk_tab_build(buf, mfirst, nm, mtab)
141 tk_tab_build(buf, vfirst, vocab, vtab)
142 let ranks: *i64 = sys_mmap((ntok + 2) * 8) as *i64
143 var k: i64 = 0
144 while k < ntok - 1 {
145 ranks[k] = tk_rank_fast(buf, mtab, tok_ptr[k] as *u8, tok_len[k], tok_ptr[k+1] as *u8, tok_len[k+1], tmp)
146 k = k + 1
147 }
148 var go: i64=1
149 while go==1 {
150 var best: i64=0-1; var besti: i64=0-1
151 k=0
152 while k<ntok-1 {
153 let r: i64=ranks[k]
154 if r>=0 { if besti<0 { best=r; besti=k } else { if r<best { best=r; besti=k } } }
155 k=k+1
156 }
157 if besti<0 { go=0 } else {
158 // MERGE besti + besti+1. Byte-level chars are laid down contiguously in `pool` and we
159 // only ever merge ADJACENT tokens, so tok_ptr[besti+1] == tok_ptr[besti]+tok_len[besti]
160 // (invariant preserved across shifts) -> the merged token is just the contiguous pool
161 // range; extend the length in place. NO per-merge sys_mmap (was leaking ~1 mapping per
162 // merge -> ~N leaked mappings/call, faulting the 2nd large call). Byte-identical output.
163 let xl: i64=tok_len[besti]; let yl: i64=tok_len[besti+1]
164 tok_len[besti]=xl+yl
165 var s: i64=besti+1; while s<ntok-1 { tok_ptr[s]=tok_ptr[s+1]; tok_len[s]=tok_len[s+1]; s=s+1 }
166 // pair-rank cache: pairs beyond besti keep their contents -> shift; the two pairs
167 // touching the merged token are the only ones whose rank changed -> recompute.
168 s=besti+1
169 while s<ntok-2 { ranks[s]=ranks[s+1]; s=s+1 }
170 ntok=ntok-1
171 if besti>0 { ranks[besti-1]=tk_rank_fast(buf, mtab, tok_ptr[besti-1] as *u8, tok_len[besti-1], tok_ptr[besti] as *u8, tok_len[besti], tmp) }
172 if besti<ntok-1 { ranks[besti]=tk_rank_fast(buf, mtab, tok_ptr[besti] as *u8, tok_len[besti], tok_ptr[besti+1] as *u8, tok_len[besti+1], tmp) }
173 }
174 }
175 i=0; while i<ntok { ids[i]=tk_tab_get(buf, vtab, tok_ptr[i] as *u8, tok_len[i]); i=i+1 }
176 sys_munmap(tmp, 512)
177 sys_munmap(pool, ilen*4 + 16)
178 sys_munmap(mtab as *u8, TK_TCAP * 16)
179 sys_munmap(vtab as *u8, TK_TCAP * 16)
180 sys_munmap(ranks as *u8, (ntok0 + 2) * 8)
181 return ntok
182}