Improve Unigram tokenization performance for long inputs

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);

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<string> {
+  *common_prefix_search(chars: string[], start = 0): Generator<string> {
     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;

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: [
+          ["<unk>", -10],
+          ["🙂", -1],
+          ["🙂a", -0.05],
+          ["a", -1],
+          ["ab", -0.3],
+          ["abc", -0.2],
+          ["b", -1],
+          ["bc", -0.4],
+          ["c", -1],
+        ],
+      },
+      "</s>",
+    );
+
+    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]> = [["<unk>", -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 }, "</s>");
+
+    const text = "abcdefghij".repeat(5000); // 50,000 characters
+    expect(unigram.encode([text])).toHaveLength(25000);
+  }, 5000); // NOTE: 5 seconds
+});