feat(thicktoken): slice tokens before decoding

Author: franklevasseur

thicktoken/src/smart-slice.test.ts

@@ -0,0 +1,77 @@
+import { SmartSlice } from './smart-slice'
+import { test, expect } from 'vitest'
+
+test('SmartSlice can slice the entire range', () => {
+  const slice = new SmartSlice([[0, 10]], 10)
+  const indices = [...slice]
+  expect(indices).toEqual([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
+})
+
+test('SmartSlice can slice multiple ranges', () => {
+  const slice = new SmartSlice(
+    [
+      [0, 3],
+      [5, 7],
+      [8, 10]
+    ],
+    10
+  )
+  const indices = [...slice]
+  expect(indices).toEqual([0, 1, 2, 5, 6, 8, 9])
+})
+
+test('SmartSlice handles empty slices', () => {
+  const slice = new SmartSlice([], 10)
+  const indices = [...slice]
+  expect(indices).toEqual([])
+})
+
+test('SmartSlice handles overlapping slices', () => {
+  const slice = new SmartSlice(
+    [
+      [0, 5],
+      [3, 8]
+    ],
+    10
+  )
+  const indices = [...slice]
+  expect(indices).toEqual([0, 1, 2, 3, 4, 5, 6, 7])
+})
+
+test('SmartSlice handles out-of-bounds slices', () => {
+  const slice = new SmartSlice(
+    [
+      [-5, 3],
+      [8, 12]
+    ],
+    10
+  )
+  const indices = [...slice]
+  expect(indices).toEqual([0, 1, 2, 8, 9])
+})
+
+test('SmartSlice handles single-point slices', () => {
+  const slice = new SmartSlice(
+    [
+      [4, 5],
+      [7, 8]
+    ],
+    10
+  )
+  const indices = [...slice]
+  expect(indices).toEqual([4, 7])
+})
+
+test('SmartSlice sorts and merges slices correctly', () => {
+  const slice = new SmartSlice(
+    [
+      [5, 7],
+      [9, 10],
+      [0, 3],
+      [2, 6]
+    ],
+    10
+  )
+  const indices = [...slice]
+  expect(indices).toEqual([0, 1, 2, 3, 4, 5, 6, 9])
+})

thicktoken/src/smart-slice.ts

@@ -0,0 +1,42 @@
+export type Slice = [number, number]
+
+export class SmartSlice {
+  private _sortedSlices: Slice[]
+
+  public constructor(slices: Slice[], total: number) {
+    // merge and sort slices here if needed
+
+    const trimmed: Slice[] = []
+    for (const [start, end] of slices) {
+      const clampedStart = Math.max(0, Math.min(total, start))
+      const clampedEnd = Math.max(0, Math.min(total, end))
+      if (clampedStart < clampedEnd) {
+        trimmed.push([clampedStart, clampedEnd])
+      }
+    }
+
+    const sorted: Slice[] = trimmed.sort((a, b) => a[0] - b[0])
+
+    const merged: Slice[] = []
+    for (const [start, end] of sorted) {
+      const last = merged[merged.length - 1]
+      if (last && start <= last[1]) {
+        // overlapping or contiguous slices, merge them
+        last[1] = Math.max(last[1], end)
+      } else {
+        // no overlap, add new slice
+        merged.push([start, end])
+      }
+    }
+
+    this._sortedSlices = merged
+  }
+
+  public *[Symbol.iterator](): Generator<number> {
+    for (const [start, end] of this._sortedSlices) {
+      for (let i = start; i < end; i++) {
+        yield i
+      }
+    }
+  }
+}

thicktoken/src/tokenizer.ts

@@ -1,11 +1,13 @@
 import cl100k_base from 'tiktoken/encoders/cl100k_base.json'
 import { Tiktoken, init } from 'tiktoken/lite/init'
 import { deepClone, mapValues, uniq } from './utils'
+import { Slice, SmartSlice } from './smart-slice'
 
 let tokenizer: TextTokenizer | null = null
 let lock: Promise<void> | false = false
 
 const CHUNK_SIZE = 100_000
+const DEFAULT_SLICES: Slice[] = [[0, Number.POSITIVE_INFINITY]]
 
 export class TextTokenizer {
   private warnOnSlowCalls = true
@@ -41,7 +43,7 @@ export class TextTokenizer {
       }
 
       const next = text.slice(i + CHUNK_SIZE - MARGIN, i + CHUNK_SIZE)
-      const unsafe = this.split(next).slice(-1)[0].length
+      const unsafe = this.split(next).slice(-1)[0]!.length
 
       yield text.slice(i, i + CHUNK_SIZE - unsafe)
 
@@ -137,17 +139,27 @@ export class TextTokenizer {
     return newObject
   }
 
-  public split(text: string): string[] {
+  public split(text: string, slices: Slice[] = DEFAULT_SLICES): string[] {
     const decoder = new TextDecoder()
+    const output = this.tokenizer.encode(text ?? '')
+
     const str: string[] = []
-    this.tokenizer.encode(text ?? '').forEach((x) => {
-      // copying to a new array because of memory allocation in WASM
-      str.push(decoder.decode(this.tokenizer.decode(new Uint32Array([x]))))
-    })
+
+    const smartSlice = new SmartSlice(slices, output.length)
+    for (const idx of smartSlice) {
+      const encodedToken = output[idx]!
+      str.push(this._decodeToken(decoder, encodedToken))
+    }
 
     return str
   }
 
+  private _decodeToken(decoder: TextDecoder, encodedToken: number): string {
+    // copying to a new array because of memory allocation in WASM
+    const copy = this.tokenizer.decode(new Uint32Array([encodedToken]))
+    return decoder.decode(copy)
+  }
+
   /**
    * Counts the number of tokens, up to a fixed ceiling, after which we return
    * The reason to have a ceiling is to avoid performance issues with very large texts

thicktoken/tsconfig.json

@@ -20,6 +20,7 @@
     "emitDecoratorMetadata": true,
     "useUnknownInCatchVariables": false,
     "forceConsistentCasingInFileNames": true,
+    "noUncheckedIndexedAccess": true,
     "outDir": "./dist",
     "rootDir": "./src",
     "types": ["node"],