249 lines
10 KiB
Plaintext
249 lines
10 KiB
Plaintext
// The tokenizer SigLIP 2's text model was trained with: Gemma's, a SentencePiece-style BPE with
|
|
// byte fallback, read from its tokenizer.json. What it does, as that file says:
|
|
//
|
|
// - spaces become "▁" (the normalizer), and the whole text is then one word;
|
|
// - each character is a piece, or when the vocabulary has no such piece, its UTF-8 bytes as
|
|
// <0xXX> pieces (byte fallback);
|
|
// - neighbouring pieces are merged, the pair with the lowest rank in the merges list first and the
|
|
// leftmost of equal ranks, until no listed pair is left;
|
|
// - <eos> goes at the end, and <pad> fills it out to max_length (64, as SigLIP 2 was trained).
|
|
//
|
|
// Before all that, text is lowercased, as SigLIP 2 was trained with lowercased text. Only ASCII
|
|
// letters are lowered here; other letters keep their case.
|
|
|
|
Gemma_Tokenizer :: struct {
|
|
text: string; // tokenizer.json, which unescaped pieces point into
|
|
vocab: Table(string, s32); // piece to id
|
|
merges: Table(u64, Merge); // left id << 32 | right id, to the merge
|
|
byte_ids: [256] s32; // <0x00>..<0xFF>
|
|
unk, eos, pad: s32;
|
|
}
|
|
|
|
Merge :: struct { rank: s32; result: s32; }
|
|
|
|
GEMMA_SPACE :: "\xE2\x96\x81"; // ▁, U+2581
|
|
|
|
// Reads tokenizer.json. Everything from the context allocator; tokenizer_free frees it.
|
|
tokenizer_load :: (path: string) -> *Gemma_Tokenizer, ok: bool, error: string {
|
|
text, read := read_entire_file(path, log_errors = false);
|
|
if !read return null, false, tprint("Cannot read %.", path);
|
|
t := New(Gemma_Tokenizer);
|
|
t.text = text;
|
|
|
|
// "vocab": { "<pad>": 0, ... }
|
|
p := find_index_from_left(text, "\"vocab\"");
|
|
if p < 0 { tokenizer_free(t); return null, false, tprint("% has no vocab.", path); }
|
|
rest := advance(text, p + 7);
|
|
if !json_expect(*rest, #char ":") || !json_expect(*rest, #char "{") { tokenizer_free(t); return null, false, tprint("% has a vocab that is not an object.", path); }
|
|
table_ensure_space(*t.vocab, 300_000);
|
|
while true {
|
|
json_skip_space(*rest);
|
|
if rest.count && rest[0] == #char "}" { advance(*rest); break; }
|
|
piece, ok := json_string(*rest);
|
|
if !ok || !json_expect(*rest, #char ":") { tokenizer_free(t); return null, false, tprint("% has a vocab entry that cannot be read.", path); }
|
|
json_skip_space(*rest);
|
|
id, number_ok, after := string_to_int(rest, T = s64);
|
|
if !number_ok { tokenizer_free(t); return null, false, tprint("% has a vocab id that cannot be read.", path); }
|
|
rest = after;
|
|
table_set(*t.vocab, piece, cast(s32) id);
|
|
json_skip_space(*rest);
|
|
if rest.count && rest[0] == #char "," advance(*rest);
|
|
}
|
|
|
|
// "merges": [ ["a", "b"], ... ]
|
|
p = find_index_from_left(rest, "\"merges\"");
|
|
if p < 0 { tokenizer_free(t); return null, false, tprint("% has no merges.", path); }
|
|
rest = advance(rest, p + 8);
|
|
if !json_expect(*rest, #char ":") || !json_expect(*rest, #char "[") { tokenizer_free(t); return null, false, tprint("% has merges that are not a list.", path); }
|
|
table_ensure_space(*t.merges, 800_000);
|
|
rank: s32 = 0;
|
|
joined: [..] u8; // each merge's two pieces end to end, to look up their result
|
|
defer array_free(joined);
|
|
while true {
|
|
json_skip_space(*rest);
|
|
if rest.count && rest[0] == #char "]" { advance(*rest); break; }
|
|
if !json_expect(*rest, #char "[") { tokenizer_free(t); return null, false, tprint("% has a merge that is not a pair (older \"a b\" merges are not read).", path); }
|
|
left, left_ok := json_string(*rest);
|
|
json_expect(*rest, #char ",");
|
|
right, right_ok := json_string(*rest);
|
|
if !left_ok || !right_ok || !json_expect(*rest, #char "]") { tokenizer_free(t); return null, false, tprint("% has a merge that cannot be read.", path); }
|
|
json_skip_space(*rest);
|
|
if rest.count && rest[0] == #char "," advance(*rest);
|
|
|
|
left_found, left_id := table_find(*t.vocab, left);
|
|
right_found, right_id := table_find(*t.vocab, right);
|
|
joined.count = 0;
|
|
array_add(*joined, ..cast([] u8) left);
|
|
array_add(*joined, ..cast([] u8) right);
|
|
result_found, result_id := table_find(*t.vocab, cast(string) joined);
|
|
if left_found && right_found && result_found {
|
|
table_set(*t.merges, ((cast(u64) left_id) << 32) | cast(u64) cast(u32) right_id, .{rank, result_id});
|
|
}
|
|
rank += 1;
|
|
}
|
|
|
|
found: bool;
|
|
found, t.unk = table_find(*t.vocab, "<unk>");
|
|
if !found { tokenizer_free(t); return null, false, tprint("% has no <unk>.", path); }
|
|
found, t.eos = table_find(*t.vocab, "<eos>");
|
|
if !found { tokenizer_free(t); return null, false, tprint("% has no <eos>.", path); }
|
|
found, t.pad = table_find(*t.vocab, "<pad>");
|
|
if !found { tokenizer_free(t); return null, false, tprint("% has no <pad>.", path); }
|
|
for 0..255 {
|
|
byte_found, id := table_find(*t.vocab, tprint("<0x%>", formatInt(it, base = 16, minimum_digits = 2)));
|
|
t.byte_ids[it] = ifx byte_found then id else t.unk;
|
|
}
|
|
return t, true, "";
|
|
}
|
|
|
|
tokenizer_free :: (t: *Gemma_Tokenizer) {
|
|
if !t return;
|
|
// Pieces with escapes were copied; the others point into text.
|
|
for _, piece: t.vocab if !tokenizer_points_into(t, piece) free(piece);
|
|
table_deinit(*t.vocab);
|
|
table_deinit(*t.merges);
|
|
free(t.text);
|
|
free(t);
|
|
}
|
|
|
|
// text as max_length token ids: its pieces (at most max_length - 1), <eos>, then <pad>. Temp.
|
|
tokenizer_encode :: (t: *Gemma_Tokenizer, text: string, max_length := 64) -> [] s64 {
|
|
// Lowercased, spaces as ▁.
|
|
normal: String_Builder;
|
|
normal.allocator = temp;
|
|
for 0..text.count - 1 {
|
|
c := text[it];
|
|
if c == #char " " append(*normal, GEMMA_SPACE);
|
|
else if c >= #char "A" && c <= #char "Z" append(*normal, c + (#char "a" - #char "A"));
|
|
else append(*normal, c);
|
|
}
|
|
word := builder_to_string(*normal,, temp);
|
|
|
|
// One piece per character; byte pieces for those the vocabulary lacks.
|
|
symbols: [..] s32;
|
|
symbols.allocator = temp;
|
|
i := 0;
|
|
while i < word.count {
|
|
length := utf8_length(word[i]);
|
|
if i + length > word.count length = word.count - i; // a cut-off sequence: its bytes
|
|
found, id := table_find(*t.vocab, slice(word, i, length));
|
|
if found array_add(*symbols, id);
|
|
else for b: 0..length - 1 array_add(*symbols, t.byte_ids[word[i + b]]);
|
|
i += length;
|
|
}
|
|
|
|
// Merges, lowest rank first, the leftmost of equal ranks.
|
|
while symbols.count > 1 {
|
|
best := -1;
|
|
best_merge: Merge;
|
|
for 0..symbols.count - 2 {
|
|
found, merge := table_find(*t.merges, ((cast(u64) symbols[it]) << 32) | cast(u64) cast(u32) symbols[it + 1]);
|
|
if found && (best < 0 || merge.rank < best_merge.rank) { best = it; best_merge = merge; }
|
|
}
|
|
if best < 0 break;
|
|
symbols[best] = best_merge.result;
|
|
array_ordered_remove_by_index(*symbols, best + 1);
|
|
}
|
|
|
|
ids := NewArray(max_length, s64,, temp);
|
|
kept := min(symbols.count, max_length - 1);
|
|
for 0..kept - 1 ids[it] = symbols[it];
|
|
ids[kept] = t.eos;
|
|
for kept + 1..max_length - 1 ids[it] = t.pad;
|
|
return ids;
|
|
}
|
|
|
|
#scope_file
|
|
|
|
// The length of the UTF-8 sequence a byte starts: 1 for ASCII and stray continuation bytes.
|
|
utf8_length :: (lead: u8) -> s64 {
|
|
if lead < 0xC0 return 1;
|
|
if lead < 0xE0 return 2;
|
|
if lead < 0xF0 return 3;
|
|
return 4;
|
|
}
|
|
|
|
tokenizer_points_into :: (t: *Gemma_Tokenizer, s: string) -> bool {
|
|
return s.data >= t.text.data && s.data < t.text.data + t.text.count;
|
|
}
|
|
|
|
json_skip_space :: (s: *string) {
|
|
while s.count && (s.*[0] == #char " " || s.*[0] == #char "\n" || s.*[0] == #char "\r" || s.*[0] == #char "\t") advance(s);
|
|
}
|
|
|
|
json_expect :: (s: *string, c: u8) -> bool {
|
|
json_skip_space(s);
|
|
if !s.count || s.*[0] != c return false;
|
|
advance(s);
|
|
return true;
|
|
}
|
|
|
|
// A JSON string at the start of s: a view into s when it has no escapes, else unescaped into a
|
|
// copy from the context allocator.
|
|
json_string :: (s: *string) -> string, ok: bool {
|
|
json_skip_space(s);
|
|
if !s.count || s.*[0] != #char "\"" return "", false;
|
|
advance(s);
|
|
end := 0;
|
|
escaped := false;
|
|
while end < s.count && s.*[end] != #char "\"" {
|
|
if s.*[end] == #char "\\" { escaped = true; end += 1; }
|
|
end += 1;
|
|
}
|
|
if end >= s.count return "", false;
|
|
raw := slice(s.*, 0, end);
|
|
advance(s, end + 1);
|
|
if !escaped return raw, true;
|
|
|
|
builder: String_Builder;
|
|
i := 0;
|
|
while i < raw.count {
|
|
c := raw[i];
|
|
if c != #char "\\" { append(*builder, c); i += 1; continue; }
|
|
if i + 1 >= raw.count return "", false;
|
|
e := raw[i + 1];
|
|
i += 2;
|
|
if e == {
|
|
case #char "n"; append(*builder, #char "\n");
|
|
case #char "t"; append(*builder, #char "\t");
|
|
case #char "r"; append(*builder, #char "\r");
|
|
case #char "b"; append(*builder, 8);
|
|
case #char "f"; append(*builder, 12);
|
|
case #char "u";
|
|
code, ok := hex4(raw, i);
|
|
if !ok return "", false;
|
|
i += 4;
|
|
if code >= 0xD800 && code < 0xDC00 && i + 6 <= raw.count && raw[i] == #char "\\" && raw[i + 1] == #char "u" {
|
|
low, low_ok := hex4(raw, i + 2);
|
|
if low_ok && low >= 0xDC00 && low < 0xE000 { code = 0x10000 + ((code - 0xD800) << 10) + (low - 0xDC00); i += 6; }
|
|
}
|
|
append_utf8(*builder, code);
|
|
case; append(*builder, e); // \" \\ \/
|
|
}
|
|
}
|
|
return builder_to_string(*builder), true;
|
|
}
|
|
|
|
hex4 :: (s: string, at: s64) -> u32, ok: bool {
|
|
if at + 4 > s.count return 0, false;
|
|
value: u32;
|
|
for 0..3 {
|
|
c := s[at + it];
|
|
digit: u32;
|
|
if c >= #char "0" && c <= #char "9" digit = c - #char "0";
|
|
else if c >= #char "a" && c <= #char "f" digit = c - #char "a" + 10;
|
|
else if c >= #char "A" && c <= #char "F" digit = c - #char "A" + 10;
|
|
else return 0, false;
|
|
value = value * 16 + digit;
|
|
}
|
|
return value, true;
|
|
}
|
|
|
|
append_utf8 :: (builder: *String_Builder, code: u32) {
|
|
if code < 0x80 { append(builder, cast(u8) code); return; }
|
|
if code < 0x800 { append(builder, cast(u8) (0xC0 | (code >> 6))); append(builder, cast(u8) (0x80 | (code & 0x3F))); return; }
|
|
if code < 0x10000 { append(builder, cast(u8) (0xE0 | (code >> 12))); append(builder, cast(u8) (0x80 | ((code >> 6) & 0x3F))); append(builder, cast(u8) (0x80 | (code & 0x3F))); return; }
|
|
append(builder, cast(u8) (0xF0 | (code >> 18))); append(builder, cast(u8) (0x80 | ((code >> 12) & 0x3F)));
|
|
append(builder, cast(u8) (0x80 | ((code >> 6) & 0x3F))); append(builder, cast(u8) (0x80 | (code & 0x3F)));
|
|
}
|