main.bend
bend-ml-bpe-tokenizer@0.1.2.0 · 656 lines · 0 definitions
1# bend-ml-bpe-tokenizer: byte-level BPE with a proved roundtrip.
2#
3# import bend-ml-bpe-tokenizer@0.1.2.0/main.bend as BPE
4#
5# A token is a raw byte (B{n}) or the result of a merge rule (M{id}).
6# The table is a list of rules from the oldest to the newest, like GPT-2's
7# merges.txt: the rule with id k joins the pair (a, b) into the token M{k}.
8# `encode` applies the rules in order; `decode` expands each token back
9# into bytes. The laws are at the end of the file, each with its proof.
10#
11# How to read the proofs: a proof in Bend is a function whose TYPE is the statement.
12# `match` does case analysis, the recursive call is the induction hypothesis,
13# `%e : P` rewrites the goal with the equality `e`, `{==}` closes when the
14# two sides are already the same term. `+x` marks a variable that may be
15# used more than once.
16
17import Base
18import bend-ml-nat-lemmas@0.1.0.0/main.bend as NL
19
20# Token: a raw byte (B) or the result of a merge rule, by rule id (M).
21type Tk is Data:
22 B{n: Nat}
23 M{k: Nat}
24
25# Merge rule: joins the pair (a, b) into the token M{id}.
26type Rule is Data:
27 Rule{id: Nat, a: Tk, b: Tk}
28
29def Tk.eq(+x: Tk, +y: Tk) -> Bool:
30 match x y:
31 case B{+m} B{+n}:
32 Nat.is_eq(m, n)
33 case M{+i} M{+j}:
34 Nat.is_eq(i, j)
35 case B{m} M{j}:
36 False{}
37 case M{i} B{n}:
38 False{}
39
40# does xs start with exactly the pair (a, b)?
41def peek(+a: Tk, +b: Tk, xs: List<&2, Tk>) -> Bool:
42 match xs:
43 case Nil{}:
44 False{}
45 case Con{+x, t}:
46 match t:
47 case Nil{}:
48 False{}
49 case Con{+y, u}:
50 Bool.and(Tk.eq(x, a), Tk.eq(y, b))
51
52# Replaces each occurrence (non-overlapping, left to right) of the pair
53# (a, b) with the token c. `hit` is always `peek(a, b, xs)`: the caller computes it,
54# because a match can only open parameters.
55def merge.go(xs: List<&2, Tk>, hit: Bool, +a: Tk, +b: Tk, +c: Tk) -> List<&2, Tk>:
56 match xs hit:
57 case Con{x, Con{y, +t}} True{}:
58 c <> merge.go(t, peek(a, b, t), a, b, c)
59 case Con{x, Nil{}} True{}:
60 x <> Nil{}
61 case Con{x, +t} False{}:
62 x <> merge.go(t, peek(a, b, t), a, b, c)
63 case Nil{} _:
64 Nil{}
65
66def merge(+a: Tk, +b: Tk, +c: Tk, +xs: List<&2, Tk>) -> List<&2, Tk>:
67 merge.go(xs, peek(a, b, xs), a, b, c)
68
69# applies the rules from the oldest to the newest; `done` keeps the ones already applied
70# (newest first), which is exactly the table that decode needs.
71def encode.go(todo: List<&2, Rule>, done: List<&2, Rule>, xs: List<&2, Tk>) -> List<&2, Tk>:
72 match todo:
73 case Nil{}:
74 xs
75 case Con{+r, rest}:
76 match r:
77 case Rule{+i, +a, +b}:
78 encode.go(rest, Rule{i, a, b} <> done, merge(a, b, M{i}, xs))
79
80# The table goes from the oldest rule to the newest (like merges.txt).
81def encode(+table: List<&2, Rule>, xs: List<&2, Tk>) -> List<&2, Tk>:
82 encode.go(table, Nil{}, xs)
83
84# is the newest rule in the table the one that creates the token tk?
85def hits(table: List<&2, Rule>, +tk: Tk) -> Bool:
86 match table:
87 case Nil{}:
88 False{}
89 case Con{Rule{+i, a, b}, rest}:
90 match tk:
91 case B{n}:
92 False{}
93 case M{+j}:
94 Nat.is_eq(i, j)
95
96# Expansion of a token into bytes. The table has the newest rule first and
97# a rule only refers to older tokens, so expanding a and b looks only at
98# the rest of the table. `hit` is always `hits(table, tk)`.
99def exp.go(table: List<&2, Rule>, hit: Bool, tk: Tk) -> List<&2, Nat>:
100 match table hit tk:
101 case Nil{} _ B{n}:
102 n <> Nil{}
103 case Nil{} _ M{j}:
104 Nil{}
105 case Con{r, rest} _ B{n}:
106 n <> Nil{}
107 case Con{Rule{i, +a, +b}, +rest} True{} M{j}:
108 List.append(&2, Nat, exp.go(rest, hits(rest, a), a), exp.go(rest, hits(rest, b), b))
109 case Con{r, +rest} False{} M{+j}:
110 exp.go(rest, hits(rest, M{j}), M{j})
111
112def exp(+table: List<&2, Rule>, +tk: Tk) -> List<&2, Nat>:
113 exp.go(table, hits(table, tk), tk)
114
115# decodes with the already reversed table (newest first)
116def dec(+table: List<&2, Rule>, ids: List<&2, Tk>) -> List<&2, Nat>:
117 match ids:
118 case Nil{}:
119 Nil{}
120 case Con{+x, t}:
121 List.append(&2, Nat, exp(table, x), dec(table, t))
122
123def decode(+table: List<&2, Rule>, ids: List<&2, Tk>) -> List<&2, Nat>:
124 dec(List.reverse(&2, Rule, table), ids)
125
126# bytes -> initial tokens
127def lift(bs: List<&2, Nat>) -> List<&2, Tk>:
128 match bs:
129 case Nil{}:
130 Nil{}
131 case Con{n, t}:
132 B{n} <> lift(t)
133
134# ---- predicates used in the laws (all computable) ----
135
136# does some rule in the table have the id i?
137def has(+D: List<&2, Rule>, +i: Nat) -> Bool:
138 match D:
139 case Nil{}:
140 False{}
141 case Con{Rule{+j, a, b}, t}:
142 Bool.or(Nat.is_eq(i, j), has(t, i))
143
144# is the token a byte, or was it created by a rule in the table?
145def kn(+D: List<&2, Rule>, tk: Tk) -> Bool:
146 match tk:
147 case B{n}:
148 True{}
149 case M{+j}:
150 has(D, j)
151
152def knl(+D: List<&2, Rule>, xs: List<&2, Tk>) -> Bool:
153 match xs:
154 case Nil{}:
155 True{}
156 case Con{x, t}:
157 Bool.and(kn(D, x), knl(D, t))
158
159# well-formed table: no id repeats
160def wf.go(todo: List<&2, Rule>, +done: List<&2, Rule>) -> Bool:
161 match todo:
162 case Nil{}:
163 True{}
164 case Con{Rule{+i, +a, +b}, rest}:
165 Bool.and(Bool.not(has(done, i)), wf.go(rest, Rule{i, a, b} <> done))
166
167def wf(table: List<&2, Rule>) -> Bool:
168 wf.go(table, Nil{})
169
170# =====================================================================
171# Training (needs no proof: it generates tables; a table is valid if wf(table))
172# =====================================================================
173# ---------------------------------------------------------------
174# training: count adjacent pairs, merge the most frequent one, repeat
175# (tie-break: the first pair that appeared, as in minbpe)
176# ---------------------------------------------------------------
177type Cnt is Data:
178 Cnt{a: Tk, b: Tk, n: Nat}
179
180# is the first Cnt of cs the pair (a, b)?
181def same_head(cs: List<&2, Cnt>, +a: Tk, +b: Tk) -> Bool:
182 match cs:
183 case Nil{}:
184 False{}
185 case Con{Cnt{+x, +y, k}, rest}:
186 Bool.and(Tk.eq(x, a), Tk.eq(y, b))
187
188# adds 1 to the counter of the pair (a, b), or creates it at the end of the list
189def bump.go(cs: List<&2, Cnt>, hit: Bool, +a: Tk, +b: Tk) -> List<&2, Cnt>:
190 match cs hit:
191 case Nil{} _:
192 Cnt{a, b, 1n} <> Nil{}
193 case Con{Cnt{x, y, +k}, rest} True{}:
194 Cnt{x, y, 1n+k} <> rest
195 case Con{c, +rest} False{}:
196 c <> bump.go(rest, same_head(rest, a, b), a, b)
197
198def bump(+cs: List<&2, Cnt>, +a: Tk, +b: Tk) -> List<&2, Cnt>:
199 bump.go(cs, same_head(cs, a, b), a, b)
200
201def count.go(xs: List<&2, Tk>, cs: List<&2, Cnt>) -> List<&2, Cnt>:
202 match xs:
203 case Nil{}:
204 cs
205 case Con{+x, +t}:
206 match t:
207 case Nil{}:
208 cs
209 case Con{+y, u}:
210 count.go(Con{y, u}, bump(cs, x, y))
211
212def pick2(better: Bool, c: Cnt, cur: Cnt) -> Cnt:
213 match better:
214 case True{}:
215 c
216 case False{}:
217 cur
218
219def pick(c: Cnt, cur: Cnt) -> Cnt:
220 match c cur:
221 case Cnt{+a, +b, +k} Cnt{+x, +y, +m}:
222 pick2(Nat.is_lt(m, k), Cnt{a, b, k}, Cnt{x, y, m})
223
224def best.go(cs: List<&2, Cnt>, cur: Cnt) -> Cnt:
225 match cs:
226 case Nil{}:
227 cur
228 case Con{c, rest}:
229 best.go(rest, pick(c, cur))
230
231def best(cs: List<&2, Cnt>) -> Maybe<&1, Cnt>:
232 match cs:
233 case Nil{}:
234 None{}
235 case Con{c, rest}:
236 Some{best.go(rest, c)}
237
238# up to n merges, ids from `next` on; m is the current best pair (best(count(xs)))
239def train.go(n: Nat, +next: Nat, +xs: List<&2, Tk>, m: Maybe<&1, Cnt>) -> List<&2, Rule>:
240 match n m:
241 case 0n _:
242 Nil{}
243 case 1n+p None{}:
244 Nil{}
245 case 1n+p Some{Cnt{+a, +b, k}}:
246 +ys = merge(a, b, M{next}, xs)
247 Rule{next, a, b} <> train.go(p, 1n+next, ys, best(count.go(ys, Nil{})))
248
249# the table from the oldest rule to the newest
250def train(n: Nat, +next: Nat, +xs: List<&2, Tk>) -> List<&2, Rule>:
251 train.go(n, next, xs, best(count.go(xs, Nil{})))
252
253# =====================================================================
254# Auxiliary lemmas for the proofs
255# =====================================================================
256
257# a type chosen by a Bool: rewriting through it refutes True == False
258def BD(b: Bool, t: Type, f: Type) -> Type:
259 match b:
260 case True{}:
261 t
262 case False{}:
263 f
264
265def nat_eq_sound(a: Nat, b: Nat, e: {True{} == Nat.is_eq(a, b) : Bool}) -> {a == b : Nat}:
266 match a b:
267 case 0n 0n:
268 {==}
269 case 0n 1n+q:
270 %e : BD(_, Unit, {0n == 1n+q : Nat})
271 Unit{}
272 case 1n+p 0n:
273 %e : BD(_, Unit, {1n+p == 0n : Nat})
274 Unit{}
275 case 1n+p 1n+q:
276 %nat_eq_sound(p, q, e) : {1n+p == 1n+_ : Nat}
277 {==}
278
279def nat_eq_refl(a: Nat) -> {True{} == Nat.is_eq(a, a) : Bool}:
280 match a:
281 case 0n:
282 {==}
283 case 1n+p:
284 nat_eq_refl(p)
285
286def and_true(p: Bool, q: Bool, e: {True{} == Bool.and(p, q) : Bool}) -> {True{} == p : Bool} & {True{} == q : Bool}:
287 match p:
288 case True{}:
289 ({==}, e)
290 case False{}:
291 %e : BD(_, Unit, {True{} == False{} : Bool} & {True{} == q : Bool})
292 Unit{}
293
294def and_intro(p: Bool, q: Bool, ep: {True{} == p : Bool}, eq: {True{} == q : Bool}) -> {True{} == Bool.and(p, q) : Bool}:
295 match p:
296 case True{}:
297 eq
298 case False{}:
299 %ep : BD(_, Unit, {True{} == False{} : Bool})
300 Unit{}
301
302def or_left(p: Bool, q: Bool, ep: {True{} == p : Bool}) -> {True{} == Bool.or(p, q) : Bool}:
303 match p:
304 case True{}:
305 {==}
306 case False{}:
307 %ep : BD(_, Unit, {True{} == q : Bool})
308 Unit{}
309
310def or_right(p: Bool, q: Bool, eq: {True{} == q : Bool}) -> {True{} == Bool.or(p, q) : Bool}:
311 match p:
312 case True{}:
313 {==}
314 case False{}:
315 eq
316
317# not p is true -> p is false
318def not_true(p: Bool, e: {True{} == Bool.not(p) : Bool}) -> {False{} == p : Bool}:
319 match p:
320 case True{}:
321 %e : BD(_, Unit, {False{} == True{} : Bool})
322 Unit{}
323 case False{}:
324 {==}
325
326def tk_eq_sound(+x: Tk, +y: Tk, e: {True{} == Tk.eq(x, y) : Bool}) -> {x == y : Tk}:
327 match x y:
328 case B{+m} B{+n}:
329 %nat_eq_sound(m, n, e) : {B{m} == B{_} : Tk}
330 {==}
331 case M{+i} M{+j}:
332 %nat_eq_sound(i, j, e) : {M{i} == M{_} : Tk}
333 {==}
334 case B{m} M{j}:
335 %e : BD(_, Unit, {B{m} == M{j} : Tk})
336 Unit{}
337 case M{i} B{n}:
338 %e : BD(_, Unit, {M{i} == B{n} : Tk})
339 Unit{}
340
341# Adding the (new) rule i on top of the table does not change the expansion of a
342# token M{j} that the old table already knows (j != i, because i is new).
343def ext_m(+D: List<&2, Rule>, +i: Nat, +a: Tk, +b: Tk, +j: Nat, h: Bool, eh: {h == Nat.is_eq(i, j) : Bool}, nf: {True{} == Bool.not(has(D, i)) : Bool}, kt: {True{} == has(D, j) : Bool}) -> {exp(Rule{i, a, b} <> D, M{j}) == exp(D, M{j}) : List<&2, Nat>}:
344 match h:
345 case True{}:
346 # h = True means i == j; but j is known (has D j) and i is not (not has D i)
347 +ij = nat_eq_sound(i, j, eh)
348 +hij = Equal.cong(Nat, Bool, k => has(D, k), i, j, ij)
349 +f1 = Equal.trans(Bool, False{}, has(D, i), has(D, j), not_true(has(D, i), nf), hij)
350 +tf = Equal.trans(Bool, True{}, has(D, j), False{}, kt, Equal.sym(Bool, False{}, has(D, j), f1))
351 %tf : BD(_, Unit, {exp.go(Rule{i, a, b} <> D, Cmp.is_eq(Nat.cmp(i, j)), M{j}) == exp.go(D, hits(D, M{j}), M{j}) : List<&2, Nat>})
352 Unit{}
353 case False{}:
354 %eh : {exp.go(Rule{i, a, b} <> D, _, M{j}) == exp(D, M{j}) : List<&2, Nat>}
355 {==}
356
357def and_true_l(p: Bool, q: Bool, e: {True{} == Bool.and(p, q) : Bool}) -> {True{} == p : Bool}:
358 match p:
359 case True{}:
360 {==}
361 case False{}:
362 %e : BD(_, Unit, {True{} == False{} : Bool})
363 Unit{}
364
365def and_true_r(p: Bool, q: Bool, e: {True{} == Bool.and(p, q) : Bool}) -> {True{} == q : Bool}:
366 match p:
367 case True{}:
368 e
369 case False{}:
370 %e : BD(_, Unit, {True{} == q : Bool})
371 Unit{}
372
373# the same for any known token (a byte never depends on the table)
374def ext_tok(+D: List<&2, Rule>, +i: Nat, +a: Tk, +b: Tk, t: Tk, kt: {True{} == kn(D, t) : Bool}, nf: {True{} == Bool.not(has(D, i)) : Bool}) -> {exp(Rule{i, a, b} <> D, t) == exp(D, t) : List<&2, Nat>}:
375 match D t:
376 case Nil{} B{n}:
377 {==}
378 case Con{Rule{j, a2, b2}, rest} B{n}:
379 {==}
380 case _ M{+j}:
381 ext_m(D, i, a, b, j, Nat.is_eq(i, j), {==}, nf, kt)
382
383def ext_dec.step(+D: List<&2, Rule>, +i: Nat, +a: Tk, +b: Tk, +x: Tk, +t: List<&2, Tk>, k1: {True{} == kn(D, x) : Bool}, nf: {True{} == Bool.not(has(D, i)) : Bool}, ih: {dec(Rule{i, a, b} <> D, t) == dec(D, t) : List<&2, Nat>}) -> {dec(Rule{i, a, b} <> D, x <> t) == dec(D, x <> t) : List<&2, Nat>}:
384 %ext_tok(D, i, a, b, x, k1, nf) : {List.append(&2, Nat, exp(Rule{i, a, b} <> D, x), dec(Rule{i, a, b} <> D, t)) == List.append(&2, Nat, _, dec(D, t)) : List<&2, Nat>}
385 %ih : {List.append(&2, Nat, exp(Rule{i, a, b} <> D, x), dec(Rule{i, a, b} <> D, t)) == List.append(&2, Nat, exp(Rule{i, a, b} <> D, x), _) : List<&2, Nat>}
386 {==}
387
388# decoding a list of known tokens does not change when the new rule is added
389def ext_dec(+D: List<&2, Rule>, +i: Nat, +a: Tk, +b: Tk, xs: List<&2, Tk>, +kx: {True{} == knl(D, xs) : Bool}, +nf: {True{} == Bool.not(has(D, i)) : Bool}) -> {dec(Rule{i, a, b} <> D, xs) == dec(D, xs) : List<&2, Nat>}:
390 match xs:
391 case Nil{}:
392 {==}
393 case Con{+x, +t}:
394 ext_dec.step(D, i, a, b, x, t, and_true_l(kn(D, x), knl(D, t), kx), nf, ext_dec(D, i, a, b, t, and_true_r(kn(D, x), knl(D, t), kx), nf))
395
396# The new token M{i} expands to the expansion of a followed by that of b
397# (looking only at the old table D).
398def exp_new(+D: List<&2, Rule>, +i: Nat, +a: Tk, +b: Tk, h: Bool, eh: {h == Nat.is_eq(i, i) : Bool}) -> {exp(Rule{i, a, b} <> D, M{i}) == List.append(&2, Nat, exp(D, a), exp(D, b)) : List<&2, Nat>}:
399 match h:
400 case True{}:
401 %eh : {exp.go(Rule{i, a, b} <> D, _, M{i}) == List.append(&2, Nat, exp(D, a), exp(D, b)) : List<&2, Nat>}
402 {==}
403 case False{}:
404 +f = Equal.trans(Bool, True{}, Nat.is_eq(i, i), False{}, nat_eq_refl(i), Equal.sym(Bool, False{}, Nat.is_eq(i, i), eh))
405 %f : BD(_, Unit, {exp(Rule{i, a, b} <> D, M{i}) == List.append(&2, Nat, exp(D, a), exp(D, b)) : List<&2, Nat>})
406 Unit{}
407
408# Central step: in place of a b t we get M{i} followed by m, where m decodes
409# like t. Both sides decode the same because M{i} expands to exp(a) ++ exp(b).
410def mp_core(+D: List<&2, Rule>, +i: Nat, +a: Tk, +b: Tk, +t: List<&2, Tk>, +m: List<&2, Tk>, ka: {True{} == kn(D, a) : Bool}, kb: {True{} == kn(D, b) : Bool}, +nf: {True{} == Bool.not(has(D, i)) : Bool}, ihm: {dec(Rule{i, a, b} <> D, m) == dec(Rule{i, a, b} <> D, t) : List<&2, Nat>}) -> {dec(Rule{i, a, b} <> D, M{i} <> m) == dec(Rule{i, a, b} <> D, a <> b <> t) : List<&2, Nat>}:
411 %Equal.sym(List<&2, Nat>, exp(Rule{i, a, b} <> D, M{i}), List.append(&2, Nat, exp(D, a), exp(D, b)), exp_new(D, i, a, b, Nat.is_eq(i, i), {==})) : {List.append(&2, Nat, _, dec(Rule{i, a, b} <> D, m)) == List.append(&2, Nat, exp(Rule{i, a, b} <> D, a), List.append(&2, Nat, exp(Rule{i, a, b} <> D, b), dec(Rule{i, a, b} <> D, t))) : List<&2, Nat>}
412 %Equal.sym(List<&2, Nat>, dec(Rule{i, a, b} <> D, m), dec(Rule{i, a, b} <> D, t), ihm) : {List.append(&2, Nat, List.append(&2, Nat, exp(D, a), exp(D, b)), _) == List.append(&2, Nat, exp(Rule{i, a, b} <> D, a), List.append(&2, Nat, exp(Rule{i, a, b} <> D, b), dec(Rule{i, a, b} <> D, t))) : List<&2, Nat>}
413 %Equal.sym(List<&2, Nat>, exp(Rule{i, a, b} <> D, a), exp(D, a), ext_tok(D, i, a, b, a, ka, nf)) : {List.append(&2, Nat, List.append(&2, Nat, exp(D, a), exp(D, b)), dec(Rule{i, a, b} <> D, t)) == List.append(&2, Nat, _, List.append(&2, Nat, exp(Rule{i, a, b} <> D, b), dec(Rule{i, a, b} <> D, t))) : List<&2, Nat>}
414 %Equal.sym(List<&2, Nat>, exp(Rule{i, a, b} <> D, b), exp(D, b), ext_tok(D, i, a, b, b, kb, nf)) : {List.append(&2, Nat, List.append(&2, Nat, exp(D, a), exp(D, b)), dec(Rule{i, a, b} <> D, t)) == List.append(&2, Nat, exp(D, a), List.append(&2, Nat, _, dec(Rule{i, a, b} <> D, t))) : List<&2, Nat>}
415 NL.append_assoc(Nat, exp(D, a), exp(D, b), dec(Rule{i, a, b} <> D, t))
416
417# if x == a and x is known, a is known too
418def kn_sub(+D: List<&2, Rule>, +x: Tk, +a: Tk, ex: {x == a : Tk}, kx: {True{} == kn(D, x) : Bool}) -> {True{} == kn(D, a) : Bool}:
419 %ex : {True{} == kn(D, _) : Bool}
420 kx
421
422# x stays out of the rewrite: if the rest decodes the same, x <> rest does too
423def mp_false(+D: List<&2, Rule>, +i: Nat, +a: Tk, +b: Tk, +x: Tk, +t: List<&2, Tk>, +mg: List<&2, Tk>, ih: {dec(Rule{i, a, b} <> D, mg) == dec(Rule{i, a, b} <> D, t) : List<&2, Nat>}) -> {dec(Rule{i, a, b} <> D, x <> mg) == dec(Rule{i, a, b} <> D, x <> t) : List<&2, Nat>}:
424 %ih : {List.append(&2, Nat, exp(Rule{i, a, b} <> D, x), dec(Rule{i, a, b} <> D, mg)) == List.append(&2, Nat, exp(Rule{i, a, b} <> D, x), _) : List<&2, Nat>}
425 {==}
426
427# the pair is x y: the tokens were equal to a and b
428def mp_hit(+D: List<&2, Rule>, +i: Nat, +a: Tk, +b: Tk, +x: Tk, +y: Tk, +t: List<&2, Tk>, +m: List<&2, Tk>, +ex: {x == a : Tk}, +ey: {y == b : Tk}, +kxx: {True{} == kn(D, x) : Bool}, +kyy: {True{} == kn(D, y) : Bool}, +nf: {True{} == Bool.not(has(D, i)) : Bool}, ihm: {dec(Rule{i, a, b} <> D, m) == dec(Rule{i, a, b} <> D, t) : List<&2, Nat>}) -> {dec(Rule{i, a, b} <> D, M{i} <> m) == dec(Rule{i, a, b} <> D, x <> y <> t) : List<&2, Nat>}:
429 %Equal.sym(Tk, x, a, ex) : {dec(Rule{i, a, b} <> D, M{i} <> m) == dec(Rule{i, a, b} <> D, _ <> y <> t) : List<&2, Nat>}
430 %Equal.sym(Tk, y, b, ey) : {dec(Rule{i, a, b} <> D, M{i} <> m) == dec(Rule{i, a, b} <> D, a <> _ <> t) : List<&2, Nat>}
431 mp_core(D, i, a, b, t, m, kn_sub(D, x, a, ex, kxx), kn_sub(D, y, b, ey, kyy), nf, ihm)
432
433# merge.go does not change what decode returns (with the new rule on top of the table)
434def mp(+D: List<&2, Rule>, +i: Nat, +a: Tk, +b: Tk, xs: List<&2, Tk>, hit: Bool, +eh: {hit == peek(a, b, xs) : Bool}, +kx: {True{} == knl(D, xs) : Bool}, +nf: {True{} == Bool.not(has(D, i)) : Bool}) -> {dec(Rule{i, a, b} <> D, merge.go(xs, hit, a, b, M{i})) == dec(Rule{i, a, b} <> D, xs) : List<&2, Nat>}:
435 match xs hit:
436 case Nil{} _:
437 {==}
438 case Con{x, Nil{}} True{}:
439 %eh : BD(_, Unit, {dec(Rule{i, a, b} <> D, merge.go(x <> Nil{}, True{}, a, b, M{i})) == dec(Rule{i, a, b} <> D, x <> Nil{}) : List<&2, Nat>})
440 Unit{}
441 case Con{+x, Con{+y, +t}} True{}:
442 mp_hit(D, i, a, b, x, y, t, merge.go(t, peek(a, b, t), a, b, M{i}), tk_eq_sound(x, a, and_true_l(Tk.eq(x, a), Tk.eq(y, b), eh)), tk_eq_sound(y, b, and_true_r(Tk.eq(x, a), Tk.eq(y, b), eh)), and_true_l(kn(D, x), Bool.and(kn(D, y), knl(D, t)), kx), and_true_l(kn(D, y), knl(D, t), and_true_r(kn(D, x), Bool.and(kn(D, y), knl(D, t)), kx)), nf, mp(D, i, a, b, t, peek(a, b, t), {==}, and_true_r(kn(D, y), knl(D, t), and_true_r(kn(D, x), Bool.and(kn(D, y), knl(D, t)), kx)), nf))
443 case Con{+x, +t} False{}:
444 mp_false(D, i, a, b, x, t, merge.go(t, peek(a, b, t), a, b, M{i}), mp(D, i, a, b, t, peek(a, b, t), {==}, and_true_r(kn(D, x), knl(D, t), kx), nf))
445
446# ---- known tokens stay known ----
447
448# adding a rule to the table only increases what is known
449def kn_mono(+D: List<&2, Rule>, +i: Nat, +a: Tk, +b: Tk, t: Tk, kt: {True{} == kn(D, t) : Bool}) -> {True{} == kn(Rule{i, a, b} <> D, t) : Bool}:
450 match t:
451 case B{n}:
452 {==}
453 case M{+j}:
454 or_right(Nat.is_eq(j, i), has(D, j), kt)
455
456# the token created by the new rule is known
457def kn_new(+D: List<&2, Rule>, +i: Nat, +a: Tk, +b: Tk) -> {True{} == kn(Rule{i, a, b} <> D, M{i}) : Bool}:
458 or_left(Nat.is_eq(i, i), has(D, i), nat_eq_refl(i))
459
460def kp(+D: List<&2, Rule>, +i: Nat, +a: Tk, +b: Tk, xs: List<&2, Tk>, hit: Bool, +eh: {hit == peek(a, b, xs) : Bool}, +kx: {True{} == knl(D, xs) : Bool}) -> {True{} == knl(Rule{i, a, b} <> D, merge.go(xs, hit, a, b, M{i})) : Bool}:
461 match xs hit:
462 case Nil{} _:
463 {==}
464 case Con{x, Nil{}} True{}:
465 %eh : BD(_, Unit, {True{} == knl(Rule{i, a, b} <> D, merge.go(x <> Nil{}, True{}, a, b, M{i})) : Bool})
466 Unit{}
467 case Con{+x, Con{+y, +t}} True{}:
468 and_intro(kn(Rule{i, a, b} <> D, M{i}), knl(Rule{i, a, b} <> D, merge.go(t, peek(a, b, t), a, b, M{i})), kn_new(D, i, a, b), kp(D, i, a, b, t, peek(a, b, t), {==}, and_true_r(kn(D, y), knl(D, t), and_true_r(kn(D, x), Bool.and(kn(D, y), knl(D, t)), kx))))
469 case Con{+x, +t} False{}:
470 and_intro(kn(Rule{i, a, b} <> D, x), knl(Rule{i, a, b} <> D, merge.go(t, peek(a, b, t), a, b, M{i})), kn_mono(D, i, a, b, x, and_true_l(kn(D, x), knl(D, t), kx)), kp(D, i, a, b, t, peek(a, b, t), {==}, and_true_r(kn(D, x), knl(D, t), kx)))
471
472# ---- induction over the rules ----
473
474# every emitted token is known to the complete (reversed) table
475def enc_known(todo: List<&2, Rule>, +done: List<&2, Rule>, +xs: List<&2, Tk>, +wfe: {True{} == wf.go(todo, done) : Bool}, +kx: {True{} == knl(done, xs) : Bool}) -> {True{} == knl(List.reverse.go(&2, Rule, todo, done), encode.go(todo, done, xs)) : Bool}:
476 match todo:
477 case Nil{}:
478 kx
479 case Con{Rule{+i, +a, +b}, +rest}:
480 enc_known(rest, Rule{i, a, b} <> done, merge.go(xs, peek(a, b, xs), a, b, M{i}), and_true_r(Bool.not(has(done, i)), wf.go(rest, Rule{i, a, b} <> done), wfe), kp(done, i, a, b, xs, peek(a, b, xs), {==}, kx))
481
482# applying all the rules does not change what decode returns
483def enc_dec(todo: List<&2, Rule>, +done: List<&2, Rule>, +xs: List<&2, Tk>, +wfe: {True{} == wf.go(todo, done) : Bool}, +kx: {True{} == knl(done, xs) : Bool}) -> {dec(List.reverse.go(&2, Rule, todo, done), encode.go(todo, done, xs)) == dec(done, xs) : List<&2, Nat>}:
484 match todo:
485 case Nil{}:
486 {==}
487 case Con{Rule{+i, +a, +b}, +rest}:
488 Equal.trans(List<&2, Nat>, dec(List.reverse.go(&2, Rule, rest, Rule{i, a, b} <> done), encode.go(rest, Rule{i, a, b} <> done, merge.go(xs, peek(a, b, xs), a, b, M{i}))), dec(Rule{i, a, b} <> done, xs), dec(done, xs), Equal.trans(List<&2, Nat>, dec(List.reverse.go(&2, Rule, rest, Rule{i, a, b} <> done), encode.go(rest, Rule{i, a, b} <> done, merge.go(xs, peek(a, b, xs), a, b, M{i}))), dec(Rule{i, a, b} <> done, merge.go(xs, peek(a, b, xs), a, b, M{i})), dec(Rule{i, a, b} <> done, xs), enc_dec(rest, Rule{i, a, b} <> done, merge.go(xs, peek(a, b, xs), a, b, M{i}), and_true_r(Bool.not(has(done, i)), wf.go(rest, Rule{i, a, b} <> done), wfe), kp(done, i, a, b, xs, peek(a, b, xs), {==}, kx)), mp(done, i, a, b, xs, peek(a, b, xs), {==}, kx, and_true_l(Bool.not(has(done, i)), wf.go(rest, Rule{i, a, b} <> done), wfe))), ext_dec(done, i, a, b, xs, kx, and_true_l(Bool.not(has(done, i)), wf.go(rest, Rule{i, a, b} <> done), wfe)))
489
490# ---- base case and final laws ----
491
492# with no rules, the bytes decode to themselves
493def lift_dec(bs: List<&2, Nat>) -> {dec(Nil{}, lift(bs)) == bs : List<&2, Nat>}:
494 match bs:
495 case Nil{}:
496 {==}
497 case Con{+n, +t}:
498 %lift_dec(t) : {n <> dec(Nil{}, lift(t)) == n <> _ : List<&2, Nat>}
499 {==}
500
501def lift_known(bs: List<&2, Nat>) -> {True{} == knl(Nil{}, lift(bs)) : Bool}:
502 match bs:
503 case Nil{}:
504 {==}
505 case Con{n, t}:
506 lift_known(t)
507
508# =====================================================================
509# LAWS (the package laws, explained in plain language in the README)
510# =====================================================================
511
512# LAW 1 (roundtrip): with a well-formed table (all rule ids
513# distinct), decoding what was encoded returns the original bytes.
514law roundtrip:
515 for +table: List<&2, Rule>
516 for +bs: List<&2, Nat>
517 for +w: {True{} == wf(table) : Bool}
518 {decode(table, encode(table, lift(bs))) == bs : List<&2, Nat>}
519
520def roundtrip(table, bs, w):
521 Equal.trans(List<&2, Nat>, decode(table, encode(table, lift(bs))), dec(Nil{}, lift(bs)), bs, enc_dec(table, Nil{}, lift(bs), w, lift_known(bs)), lift_dec(bs))
522
523# LAW 2 (vocabulary): every token emitted by encode is a raw byte or was
524# created by one of the table's rules; an unknown id never appears.
525law vocab_bound:
526 for +table: List<&2, Rule>
527 for +bs: List<&2, Nat>
528 for +w: {True{} == wf(table) : Bool}
529 {True{} == knl(List.reverse(&2, Rule, table), encode(table, lift(bs))) : Bool}
530
531def vocab_bound(table, bs, w):
532 enc_known(table, Nil{}, lift(bs), w, lift_known(bs))
533
534# LAW 3 (concatenation): decoding two token lists together is decoding
535# each one and joining the bytes.
536law dec_append:
537 for +D: List<&2, Rule>
538 for xs: List<&2, Tk>
539 for +ys: List<&2, Tk>
540 {dec(D, List.append(&2, Tk, xs, ys)) == List.append(&2, Nat, dec(D, xs), dec(D, ys)) : List<&2, Nat>}
541
542def dec_append(D, xs, ys):
543 match xs:
544 case Nil{}:
545 {==}
546 case Con{+x, +t}:
547 %Equal.sym(List<&2, Nat>, dec(D, List.append(&2, Tk, t, ys)), List.append(&2, Nat, dec(D, t), dec(D, ys)), dec_append(D, t, ys)) : {List.append(&2, Nat, exp(D, x), _) == List.append(&2, Nat, List.append(&2, Nat, exp(D, x), dec(D, t)), dec(D, ys)) : List<&2, Nat>}
548 Equal.sym(List<&2, Nat>, List.append(&2, Nat, List.append(&2, Nat, exp(D, x), dec(D, t)), dec(D, ys)), List.append(&2, Nat, exp(D, x), List.append(&2, Nat, dec(D, t), dec(D, ys))), NL.append_assoc(Nat, exp(D, x), dec(D, t), dec(D, ys)))
549
550# =====================================================================
551# LAW 4: train generates well-formed tables (closes the roundtrip gap)
552# =====================================================================
553#
554# train creates rules with ids next, next+1, next+2, ... None repeats, so the
555# table satisfies wf, which is the hypothesis of the roundtrip law. Proof by induction
556# over train.go, keeping the invariant "all ids already used are smaller
557# than the next id".
558
559# n < 1 + n
560def lt_succ(n: Nat) -> {True{} == Nat.is_lt(n, 1n+n) : Bool}:
561 match n:
562 case 0n:
563 {==}
564 case 1n+p:
565 lt_succ(p)
566
567# j < i -> i != j
568def lt_ne(j: Nat, i: Nat, e: {True{} == Nat.is_lt(j, i) : Bool}) -> {False{} == Nat.is_eq(i, j) : Bool}:
569 match j i:
570 case 0n 0n:
571 %e : BD(_, Unit, {False{} == Nat.is_eq(0n, 0n) : Bool})
572 Unit{}
573 case 0n 1n+q:
574 {==}
575 case 1n+p 0n:
576 %e : BD(_, Unit, {False{} == Nat.is_eq(0n, 1n+p) : Bool})
577 Unit{}
578 case 1n+p 1n+q:
579 lt_ne(p, q, e)
580
581# j < i -> j < 1 + i
582def lt_up(j: Nat, i: Nat, e: {True{} == Nat.is_lt(j, i) : Bool}) -> {True{} == Nat.is_lt(j, 1n+i) : Bool}:
583 match j i:
584 case 0n 0n:
585 %e : BD(_, Unit, {True{} == Nat.is_lt(0n, 1n) : Bool})
586 Unit{}
587 case 0n 1n+q:
588 {==}
589 case 1n+p 0n:
590 %e : BD(_, Unit, {True{} == Nat.is_lt(1n+p, 1n) : Bool})
591 Unit{}
592 case 1n+p 1n+q:
593 lt_up(p, q, e)
594
595# all ids in the rule list are smaller than i
596def below(D: List<&2, Rule>, +i: Nat) -> Bool:
597 match D:
598 case Nil{}:
599 True{}
600 case Con{Rule{+j, a, b}, t}:
601 Bool.and(Nat.is_lt(j, i), below(t, i))
602
603def below_has(D: List<&2, Rule>, +i: Nat, +e: {True{} == below(D, i) : Bool}) -> {False{} == has(D, i) : Bool}:
604 match D:
605 case Nil{}:
606 {==}
607 case Con{Rule{+j, a, b}, +t}:
608 %lt_ne(j, i, and_true_l(Nat.is_lt(j, i), below(t, i), e)) : {False{} == Bool.or(_, has(t, i)) : Bool}
609 below_has(t, i, and_true_r(Nat.is_lt(j, i), below(t, i), e))
610
611def below_up(D: List<&2, Rule>, +i: Nat, +e: {True{} == below(D, i) : Bool}) -> {True{} == below(D, 1n+i) : Bool}:
612 match D:
613 case Nil{}:
614 {==}
615 case Con{Rule{+j, a, b}, +t}:
616 and_intro(Nat.is_lt(j, 1n+i), below(t, 1n+i), lt_up(j, i, and_true_l(Nat.is_lt(j, i), below(t, i), e)), below_up(t, i, and_true_r(Nat.is_lt(j, i), below(t, i), e)))
617
618# adding the rule with id `next` (greater than all those in `done`) keeps the table "below 1+next"
619def below_cons(+a: Tk, +b: Tk, +D: List<&2, Rule>, +i: Nat, +e: {True{} == below(D, i) : Bool}) -> {True{} == below(Rule{i, a, b} <> D, 1n+i) : Bool}:
620 and_intro(Nat.is_lt(i, 1n+i), below(D, 1n+i), lt_succ(i), below_up(D, i, e))
621
622# not (has D i) when all the ids in D are smaller than i
623def below_fresh(+D: List<&2, Rule>, +i: Nat, +e: {True{} == below(D, i) : Bool}) -> {True{} == Bool.not(has(D, i)) : Bool}:
624 %below_has(D, i, e) : {True{} == Bool.not(_) : Bool}
625 {==}
626
627# train.go generates rules with ids next, next+1, ...: none repeats, whatever the corpus
628def tw(n: Nat, +next: Nat, +xs: List<&2, Tk>, m: Maybe<&1, Cnt>, +done: List<&2, Rule>, +bl: {True{} == below(done, next) : Bool}) -> {True{} == wf.go(train.go(n, next, xs, m), done) : Bool}:
629 match n m:
630 case 0n _:
631 {==}
632 case 1n++p None{}:
633 {==}
634 case 1n++p Some{Cnt{+a, +b, k}}:
635 and_intro(Bool.not(has(done, next)), wf.go(train.go(p, 1n+next, merge(a, b, M{next}, xs), best(count.go(merge(a, b, M{next}, xs), Nil{}))), Rule{next, a, b} <> done), below_fresh(done, next, bl), tw(p, 1n+next, merge(a, b, M{next}, xs), best(count.go(merge(a, b, M{next}, xs), Nil{})), Rule{next, a, b} <> done, below_cons(a, b, done, next, bl)))
636
637# LAW (train_wf): the table that train returns is always well formed, so the roundtrip holds for it
638law train_wf:
639 for n: Nat
640 for +next: Nat
641 for +xs: List<&2, Tk>
642 {True{} == wf(train(n, next, xs)) : Bool}
643
644def train_wf(n, next, xs):
645 tw(n, next, xs, best(count.go(xs, Nil{})), Nil{}, {==})
646
647# LAW 5 (roundtrip_trained): train a table on any corpus, with as many merges as you like:
648# decoding what it encodes always returns the original bytes. No hypothesis at all.
649law roundtrip_trained:
650 for +n: Nat
651 for +corpus: List<&2, Nat>
652 for +bs: List<&2, Nat>
653 {decode(train(n, 0n, lift(corpus)), encode(train(n, 0n, lift(corpus)), lift(bs))) == bs : List<&2, Nat>}
654
655def roundtrip_trained(n, corpus, bs):
656 roundtrip(train(n, 0n, lift(corpus)), bs, train_wf(n, 0n, lift(corpus)))