nx_wordpiece_lib.nx source
↩ module page · 244 lines · 10016 B
1// nx_wordpiece_lib.nx -- SOVEREIGN WordPiece tokenizer over a BERT uncased vocab.txt, a library so the cross-encoder
2// arm of nx_beir_eval and its gate drive it in-process (search rung R0, 2026-09-14). ABSENT-PROVEN before writing
3// (nx_absent wordpiece buildroot/runtime nx: matches=0 coverage_complete=1). The vocab is read whole through
4// sys_read_file (no cap), every token is hashed into an open-addressing table sized from the token count, and every
5// hit is verified by bytes, so a hash collision can never return the wrong id. Basic tokenisation follows the HF
6// BertTokenizer with do_lower_case: ASCII lowercase, every byte below 33 splits, every ASCII punctuation byte is its
7// own token, a word longer than WP_MAX_WORD_CHARS is [UNK] whole, then greedy longest-match wordpiece with the
8// two-byte continuation prefix. DEVIATION, stated: bytes above 0x7f are kept as word characters (no NFD accent
9// stripping, no CJK isolation), so an accented word normally becomes [UNK] pieces. license_tier: ORIGINAL
10import "nx_syscalls.nx"
11
12// g-box slots; the caller mmaps WP_G_BYTES
13const WP_G_VBUF: i64 = 0 // vocab.txt bytes
14const WP_G_VLEN: i64 = 1
15const WP_G_NTOK: i64 = 2
16const WP_G_TOFF: i64 = 3 // *i64 token start offset per id
17const WP_G_TLEN: i64 = 4 // *i64 token byte length per id
18const WP_G_HT: i64 = 5 // *i64 open-addressing table, 2 slots per bucket: hash, id+1 (0 = empty)
19const WP_G_HCAP: i64 = 6 // bucket count, a power of two derived from the token count
20const WP_G_UNK: i64 = 7
21const WP_G_CLS: i64 = 8
22const WP_G_SEP: i64 = 9
23const WP_G_PAD: i64 = 10
24const WP_G_LOADED: i64 = 11
25const WP_G_WSCR: i64 = 12 // *u8 lowercased-word scratch, WP_MAX_WORD_CHARS bytes
26const WP_G_PSCR: i64 = 13 // *i64 piece scratch, WP_MAX_WORD_CHARS slots
27const WP_G_SLOTS: i64 = 14
28const WP_G_BYTES: i64 = 112 // WP_G_SLOTS * 8
29
30// HF BertTokenizer max_input_chars_per_word: a longer word is [UNK] as a whole (the reference implementation's bound)
31const WP_MAX_WORD_CHARS: i64 = 100
32// the continuation prefix is two of this byte; the nx lexer refuses the character inside a string literal
33const WP_HASH_BYTE: i64 = 35
34// djb2 string hash: seed and multiplier are the algorithm's published parameters, not tunables
35const WP_DJB2_SEED: i64 = 5381
36const WP_DJB2_MUL: i64 = 33
37const WP_LF: i64 = 10
38const WP_CR: i64 = 13
39const WP_SPACE: i64 = 32
40const WP_UPPER_A: i64 = 65
41const WP_UPPER_Z: i64 = 90
42const WP_CASE_DELTA: i64 = 32
43const WP_DEL: i64 = 127
44
45func wp_hash_run(h0: i64, buf: *u8, off: i64, len: i64) -> i64 {
46 var h: i64 = h0
47 var i: i64 = 0
48 while i < len { h = h * WP_DJB2_MUL + (buf[off + i] as i64); i = i + 1 }
49 return h
50}
51
52// 1 if vocab token id is exactly buf[off..off+len) with the continuation prefix when pre==1
53func wp_tok_eq(g: *i64, id: i64, pre: i64, buf: *u8, off: i64, len: i64) -> i64 {
54 let toff: *i64 = g[WP_G_TOFF] as *i64
55 let tlen: *i64 = g[WP_G_TLEN] as *i64
56 let vb: *u8 = g[WP_G_VBUF] as *u8
57 var need: i64 = len
58 if pre == 1 { need = len + 2 }
59 if tlen[id] != need { return 0 }
60 var p: i64 = toff[id]
61 if pre == 1 {
62 if vb[p] != (WP_HASH_BYTE as u8) { return 0 }
63 if vb[p + 1] != (WP_HASH_BYTE as u8) { return 0 }
64 p = p + 2
65 }
66 var i: i64 = 0
67 while i < len { if vb[p + i] != buf[off + i] { return 0 } i = i + 1 }
68 return 1
69}
70
71// id of buf[off..off+len) (with the continuation prefix when pre==1), or -1
72func wp_lookup(g: *i64, pre: i64, buf: *u8, off: i64, len: i64) -> i64 {
73 var h: i64 = WP_DJB2_SEED
74 if pre == 1 { h = h * WP_DJB2_MUL + WP_HASH_BYTE; h = h * WP_DJB2_MUL + WP_HASH_BYTE }
75 h = wp_hash_run(h, buf, off, len)
76 let ht: *i64 = g[WP_G_HT] as *i64
77 let cap: i64 = g[WP_G_HCAP]
78 var idx: i64 = h & (cap - 1)
79 var probes: i64 = 0
80 while probes < cap {
81 let stored: i64 = ht[idx * 2 + 1]
82 if stored == 0 { return 0 - 1 }
83 if ht[idx * 2] == h { if wp_tok_eq(g, stored - 1, pre, buf, off, len) == 1 { return stored - 1 } }
84 idx = (idx + 1) & (cap - 1)
85 probes = probes + 1
86 }
87 return 0 - 1
88}
89
90// load vocab.txt (line index = id). Returns 1 loaded, 0 refused (file absent or a special token missing).
91func wp_load(g: *i64, path: *u8) -> i64 {
92 g[WP_G_LOADED] = 0
93 let lb: *i64 = sys_mmap(8) as *i64
94 let vb: *u8 = sys_read_file(path, lb)
95 let vlen: i64 = lb[0]
96 if vlen <= 0 { return 0 }
97 var n: i64 = 0
98 var i: i64 = 0
99 while i < vlen { if vb[i] == (WP_LF as u8) { n = n + 1 } i = i + 1 }
100 if vb[vlen - 1] != (WP_LF as u8) { n = n + 1 }
101 let toff: *i64 = sys_mmap(n * 8) as *i64
102 let tlen: *i64 = sys_mmap(n * 8) as *i64
103 var id: i64 = 0
104 var s: i64 = 0
105 i = 0
106 while i <= vlen {
107 var atend: i64 = 0
108 if i == vlen { atend = 1 } else { if vb[i] == (WP_LF as u8) { atend = 1 } }
109 if atend == 1 {
110 var e: i64 = i
111 if e > s { if vb[e - 1] == (WP_CR as u8) { e = e - 1 } }
112 if id < n { toff[id] = s; tlen[id] = e - s; id = id + 1 }
113 s = i + 1
114 }
115 i = i + 1
116 }
117 // load factor at most one half: the smallest power of two above twice the token count
118 var cap: i64 = 1
119 while cap < n * 2 { cap = cap * 2 }
120 let ht: *i64 = sys_mmap(cap * 16) as *i64
121 i = 0
122 while i < cap * 2 { ht[i] = 0; i = i + 1 }
123 g[WP_G_VBUF] = vb as i64
124 g[WP_G_VLEN] = vlen
125 g[WP_G_NTOK] = n
126 g[WP_G_TOFF] = toff as i64
127 g[WP_G_TLEN] = tlen as i64
128 g[WP_G_HT] = ht as i64
129 g[WP_G_HCAP] = cap
130 id = 0
131 while id < n {
132 let h: i64 = wp_hash_run(WP_DJB2_SEED, vb, toff[id], tlen[id])
133 var idx: i64 = h & (cap - 1)
134 while ht[idx * 2 + 1] != 0 { idx = (idx + 1) & (cap - 1) }
135 ht[idx * 2] = h
136 ht[idx * 2 + 1] = id + 1
137 id = id + 1
138 }
139 g[WP_G_WSCR] = sys_mmap(WP_MAX_WORD_CHARS) as i64
140 g[WP_G_PSCR] = sys_mmap(WP_MAX_WORD_CHARS * 8) as i64
141 g[WP_G_UNK] = wp_lookup(g, 0, "[UNK]" as *u8, 0, 5)
142 g[WP_G_CLS] = wp_lookup(g, 0, "[CLS]" as *u8, 0, 5)
143 g[WP_G_SEP] = wp_lookup(g, 0, "[SEP]" as *u8, 0, 5)
144 g[WP_G_PAD] = wp_lookup(g, 0, "[PAD]" as *u8, 0, 5)
145 if g[WP_G_UNK] < 0 { return 0 }
146 if g[WP_G_CLS] < 0 { return 0 }
147 if g[WP_G_SEP] < 0 { return 0 }
148 if g[WP_G_PAD] < 0 { return 0 }
149 g[WP_G_LOADED] = 1
150 return 1
151}
152
153func wp_is_punct(c: i64) -> i64 {
154 if c >= 33 { if c <= 47 { return 1 } }
155 if c >= 58 { if c <= 64 { return 1 } }
156 if c >= 91 { if c <= 96 { return 1 } }
157 if c >= 123 { if c <= 126 { return 1 } }
158 return 0
159}
160
161// greedy longest-match wordpieces of one lowercased word buf[off..off+len) appended at out[n]; returns the new n.
162// A word with an unmatchable remainder is [UNK] whole, exactly as the reference does.
163func wp_word(g: *i64, buf: *u8, off: i64, len: i64, out: *i64, n: i64) -> i64 {
164 if len > WP_MAX_WORD_CHARS { out[n] = g[WP_G_UNK]; return n + 1 }
165 let tmp: *i64 = g[WP_G_PSCR] as *i64
166 var np: i64 = 0
167 var ok: i64 = 1
168 var start: i64 = 0
169 while start < len {
170 var end: i64 = len
171 var found: i64 = 0 - 1
172 var fend: i64 = 0
173 var go: i64 = 1
174 while go == 1 {
175 if end <= start { go = 0 } else {
176 var pre: i64 = 0
177 if start > 0 { pre = 1 }
178 let id: i64 = wp_lookup(g, pre, buf, off + start, end - start)
179 if id >= 0 { found = id; fend = end; go = 0 } else { end = end - 1 }
180 }
181 }
182 if found < 0 { ok = 0; start = len } else { tmp[np] = found; np = np + 1; start = fend }
183 }
184 if ok == 0 { out[n] = g[WP_G_UNK]; return n + 1 }
185 var i: i64 = 0
186 while i < np { out[n + i] = tmp[i]; i = i + 1 }
187 return n + np
188}
189
190// tokenise text[off..off+len) into out; returns the token count. out needs at most len slots: a word of k bytes
191// yields at most k pieces and a punctuation byte yields exactly one token, so the bound is the byte count.
192func wp_tokenize(g: *i64, text: *u8, off: i64, len: i64, out: *i64) -> i64 {
193 let scr: *u8 = g[WP_G_WSCR] as *u8
194 var n: i64 = 0
195 var wl: i64 = 0
196 var i: i64 = off
197 let end: i64 = off + len
198 while i <= end {
199 var c: i64 = WP_SPACE
200 if i < end { c = text[i] as i64 }
201 var kind: i64 = 0 // 0 word byte, 1 separator, 2 punctuation
202 if c < 33 { kind = 1 }
203 if c == WP_DEL { kind = 1 }
204 if wp_is_punct(c) == 1 { kind = 2 }
205 if kind == 0 {
206 var lc: i64 = c
207 if c >= WP_UPPER_A { if c <= WP_UPPER_Z { lc = c + WP_CASE_DELTA } }
208 if wl < WP_MAX_WORD_CHARS { scr[wl] = lc as u8 }
209 wl = wl + 1
210 } else {
211 if wl > 0 { n = wp_word(g, scr, 0, wl, out, n); wl = 0 }
212 if kind == 2 {
213 scr[0] = c as u8
214 let pid: i64 = wp_lookup(g, 0, scr, 0, 1)
215 if pid >= 0 { out[n] = pid } else { out[n] = g[WP_G_UNK] }
216 n = n + 1
217 }
218 }
219 i = i + 1
220 }
221 return n
222}
223
224// [CLS] q [SEP] d [SEP] with token types 0 over the query segment and 1 over the document segment, truncating the
225// document to fit maxlen (the query is kept whole, as a query is always the shorter side here; a query alone longer
226// than maxlen is cut to fit). Returns the sequence length T.
227func wp_pair(g: *i64, qtok: *i64, nq0: i64, dtok: *i64, nd0: i64, maxlen: i64, ids: *i64, types: *i64) -> i64 {
228 var nq: i64 = nq0
229 var nd: i64 = nd0
230 if nq > maxlen - 3 { nq = maxlen - 3 }
231 if nq < 0 { nq = 0 }
232 var room: i64 = maxlen - 3 - nq
233 if room < 0 { room = 0 }
234 if nd > room { nd = room }
235 var t: i64 = 0
236 ids[t] = g[WP_G_CLS]; types[t] = 0; t = t + 1
237 var i: i64 = 0
238 while i < nq { ids[t] = qtok[i]; types[t] = 0; t = t + 1; i = i + 1 }
239 ids[t] = g[WP_G_SEP]; types[t] = 0; t = t + 1
240 i = 0
241 while i < nd { ids[t] = dtok[i]; types[t] = 1; t = t + 1; i = i + 1 }
242 ids[t] = g[WP_G_SEP]; types[t] = 1; t = t + 1
243 return t
244}