diff --git a/src/core/tokenizerModelImplementations/Unigram.ts b/src/core/tokenizerModelImplementations/Unigram.ts index 6c6e70f..55c0df3 100644 --- a/src/core/tokenizerModelImplementations/Unigram.ts +++ b/src/core/tokenizerModelImplementations/Unigram.ts @@ -70,11 +70,8 @@ class Unigram extends TokenizerModel { while (begin_pos < chars.length) { let has_single_node = false; - const tokens: string[] = []; - const sliced = chars.slice(begin_pos).join(""); - const prefixed_tokens = this.trie.common_prefix_search(sliced); + const prefixed_tokens = this.trie.common_prefix_search(chars, begin_pos); for (const token of prefixed_tokens) { - tokens.push(token); const token_id = this.tokens_to_ids.get(token)!; const token_score = this.scores[token_id]; const n = len(token); diff --git a/src/utils/data-structures/CharTrie.ts b/src/utils/data-structures/CharTrie.ts index d1f654b..902e68f 100644 --- a/src/utils/data-structures/CharTrie.ts +++ b/src/utils/data-structures/CharTrie.ts @@ -62,16 +62,18 @@ class CharTrie { } /** - * Searches the trie for all strings with a common prefix of `text`. - * @param text The common prefix to search for. - * @yields Each string in the trie that has `text` as a prefix. + * Searches the trie for stored strings that match `chars` starting at `start`. + * @param chars The input characters to search. + * @param start The index to start searching from. + * @yields Each stored string that is a prefix of `chars` starting at `start`. */ - *common_prefix_search(text: string): Generator { + *common_prefix_search(chars: string[], start = 0): Generator { let node: CharTrieNode | undefined = this.root; if (node === undefined) return; let prefix = ""; - for (const ch of text) { + for (let i = start; i < chars.length; ++i) { + const ch = chars[i]; prefix += ch; node = node.children.get(ch); if (node === undefined) return; diff --git a/tests/unigram.test.ts b/tests/unigram.test.ts new file mode 100644 index 0000000..fd07d0a --- /dev/null +++ b/tests/unigram.test.ts @@ -0,0 +1,57 @@ +import { Unigram } from "../src"; +import CharTrie from "../src/utils/data-structures/CharTrie"; + +describe("CharTrie", () => { + it("searches the trie from a character offset", () => { + const trie = new CharTrie(); + trie.extend(["🙂", "🙂a", "a"]); + + const chars = Array.from("x🙂ab"); + expect([...trie.common_prefix_search(chars, 1)]).toEqual(["🙂", "🙂a"]); + expect([...trie.common_prefix_search(chars, 0)]).toEqual([]); + expect([...trie.common_prefix_search(chars, chars.length)]).toEqual([]); + expect([...trie.common_prefix_search(chars, chars.length + 1)]).toEqual([]); + }); +}); + +describe("Unigram", () => { + it("selects the best-scoring path with overlapping Unicode tokens", () => { + const unigram = new Unigram( + { + type: "Unigram", + unk_id: 0, + vocab: [ + ["", -10], + ["🙂", -1], + ["🙂a", -0.05], + ["a", -1], + ["ab", -0.3], + ["abc", -0.2], + ["b", -1], + ["bc", -0.4], + ["c", -1], + ], + }, + "", + ); + + expect(unigram.encode(["🙂abc🙂a"])).toEqual(["🙂a", "bc", "🙂a"]); + expect(unigram.encode(["x"])).toEqual(["x"]); + }); + + it("tokenizes long inputs without quadratic suffix materialization", () => { + const vocab: Array<[string, number]> = [["", -10]]; + for (const c of "abcdefghij") { + vocab.push([c, -1]); + } + for (const c1 of "abcdefghij") { + for (const c2 of "abcdefghij") { + vocab.push([`${c1}${c2}`, -0.5]); + } + } + const unigram = new Unigram({ type: "Unigram", unk_id: 0, vocab }, ""); + + const text = "abcdefghij".repeat(5000); // 50,000 characters + expect(unigram.encode([text])).toHaveLength(25000); + }, 5000); // NOTE: 5 seconds +});