gpt-tokenizer 2.7.0 → 2.8.0

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "gpt-tokenizer",
-  "version": "2.7.0",
+  "version": "2.8.0",
   "description": "A pure JavaScript implementation of a BPE tokenizer (Encoder/Decoder) for GPT-2 / GPT-3 / GPT-4 and other OpenAI models",
   "keywords": [
     "BPE",

package/src/BytePairEncodingCore.ts

@@ -1,5 +1,6 @@
 /* eslint-disable no-continue */
 
+import { DEFAULT_MERGE_CACHE_SIZE } from './constants.js'
 import { compareUint8Arrays, isAscii, tryConvertToString } from './utfUtil.js'
 import { escapeRegExp } from './util.js'
 
@@ -9,6 +10,13 @@ export interface BytePairEncodingConfig {
   bytePairRankDecoder: RawBytePairRanks
   specialTokensEncoder?: Map<string, number>
   tokenSplitRegex: RegExp
+  /**
+   * LRU cache for merged tokens pairs.
+   * Increasing this value should make encoding similar strings faster,
+   * but will consume more memory.
+   * @default 100000
+   */
+  mergeCacheSize?: number
 }
 
 const emptyBuffer = new Uint8Array(0)
@@ -36,14 +44,21 @@ export class BytePairEncodingCore {
   private specialTokensDecoder: Map<number, string>
   private specialTokenPatternRegex: RegExp
   private textEncoder = new TextEncoder()
+  private mergeCache?: Map<string, number[]>
+  private mergeCacheSize: number
 
   constructor({
     bytePairRankDecoder,
     specialTokensEncoder,
     tokenSplitRegex,
+    mergeCacheSize = DEFAULT_MERGE_CACHE_SIZE,
   }: BytePairEncodingConfig) {
     this.bytePairRankDecoder = bytePairRankDecoder
     this.bytePairStringRankEncoder = new Map<string, number>()
+    this.mergeCacheSize = mergeCacheSize
+    if (mergeCacheSize > 0) {
+      this.mergeCache = new Map()
+    }
 
     // size without array holes (which may be present in the encoder)
     this.mergeableBytePairRankCount = Object.keys(bytePairRankDecoder).length
@@ -73,12 +88,22 @@ export class BytePairEncodingCore {
     )
     const allSpecialTokensRegex = escapedSpecialTokens.join('|')
     try {
-      this.specialTokenPatternRegex = new RegExp(allSpecialTokensRegex)
+      this.specialTokenPatternRegex = new RegExp(allSpecialTokensRegex, 'y')
     } catch {
       throw new Error('Invalid regular expression pattern.')
     }
   }
 
+  setMergeCacheSize(newSize: number): void {
+    if (this.mergeCacheSize === 0 && newSize > 0) {
+      this.mergeCache = new Map()
+    }
+    this.mergeCacheSize = newSize
+    if (newSize === 0) {
+      this.mergeCache = undefined
+    }
+  }
+
   *encodeNativeGenerator(
     text: string,
     allowedSpecial?: Set<string>,
@@ -96,7 +121,10 @@ export class BytePairEncodingCore {
 
       const endIndex = nextSpecialStartIndex ?? text.length
 
-      const textBeforeSpecial = text.slice(startIndex, endIndex)
+      const textBeforeSpecial =
+        startIndex === 0 && endIndex === text.length
+          ? text
+          : text.slice(startIndex, endIndex)
 
       for (const [match] of textBeforeSpecial.matchAll(this.tokenSplitRegex)) {
         const token = this.getBpeRankFromString(match)
@@ -146,7 +174,10 @@ export class BytePairEncodingCore {
 
       const endIndex = nextSpecialStartIndex ?? text.length
 
-      const textBeforeSpecial = text.slice(startIndex, endIndex)
+      const textBeforeSpecial =
+        startIndex === 0 && endIndex === text.length
+          ? text
+          : text.slice(startIndex, endIndex)
 
       for (const [match] of textBeforeSpecial.matchAll(this.tokenSplitRegex)) {
         const token = this.getBpeRankFromString(match)
@@ -178,6 +209,56 @@ export class BytePairEncodingCore {
     return tokensArray
   }
 
+  countNative(text: string, allowedSpecial?: Set<string>): number {
+    let startIndex = 0
+    let tokensCount = 0
+
+    // eslint-disable-next-line no-constant-condition
+    while (true) {
+      const nextSpecialMatch = this.findNextSpecialToken(
+        text,
+        allowedSpecial,
+        startIndex,
+      )
+      const nextSpecialStartIndex = nextSpecialMatch?.[0]
+
+      const endIndex = nextSpecialStartIndex ?? text.length
+
+      const textBeforeSpecial =
+        startIndex === 0 && endIndex === text.length
+          ? text
+          : text.slice(startIndex, endIndex)
+
+      for (const [match] of textBeforeSpecial.matchAll(this.tokenSplitRegex)) {
+        const token = this.getBpeRankFromString(match)
+        if (token !== undefined) {
+          tokensCount++
+
+          continue
+        }
+
+        const tokens = this.bytePairEncode(match)
+        tokensCount += tokens.length
+      }
+
+      if (nextSpecialStartIndex !== undefined) {
+        const specialToken = nextSpecialMatch![1]
+        const specialTokenValue = this.specialTokensEncoder.get(specialToken)
+        if (specialTokenValue === undefined) {
+          throw new Error(
+            `Special token "${specialToken}" is not in the special token encoder.`,
+          )
+        }
+        tokensCount++
+        startIndex = nextSpecialStartIndex + specialToken.length
+      } else {
+        break
+      }
+    }
+
+    return tokensCount
+  }
+
   *decodeNativeGenerator(
     tokens: Iterable<number>,
   ): Generator<Uint8Array | string, void, void> {
@@ -308,9 +389,8 @@ export class BytePairEncodingCore {
 
     // eslint-disable-next-line no-constant-condition
     while (true) {
-      const nextSpecialMatch = this.specialTokenPatternRegex.exec(
-        text.slice(Math.max(0, searchIndex)),
-      )
+      this.specialTokenPatternRegex.lastIndex = searchIndex
+      const nextSpecialMatch = this.specialTokenPatternRegex.exec(text)
 
       if (!nextSpecialMatch) {
         return undefined
@@ -341,14 +421,34 @@ export class BytePairEncodingCore {
     return this.specialTokensDecoder.get(tokenRank)
   }
 
+  private addToMergeCache(key: string, value: number[]): void {
+    if (!this.mergeCache) return
+
+    if (this.mergeCache.size >= this.mergeCacheSize) {
+      // Remove least recently used item (first item)
+      const firstKey = this.mergeCache.keys().next().value!
+      this.mergeCache.delete(firstKey)
+    }
+    this.mergeCache.set(key, value)
+  }
+
   private bytePairEncode(input: string): number[] {
     if (input.length === 1 && isAscii(input.codePointAt(0)!)) {
       return [this.getBpeRankFromStringOrThrow(input)]
     }
 
-    const inputBytes = this.textEncoder.encode(input)
+    if (this.mergeCache?.has(input)) {
+      const result = this.mergeCache.get(input)!
+      // Move to end to mark as recently used
+      this.mergeCache.delete(input)
+      this.mergeCache.set(input, result)
+      return result
+    }
 
-    return this.bytePairMerge(inputBytes)
+    const inputBytes = this.textEncoder.encode(input)
+    const result = this.bytePairMerge(inputBytes)
+    this.addToMergeCache(input, result)
+    return result
   }
 
   private bytePairMerge(

package/src/GptEncoding.ts

@@ -12,10 +12,8 @@ import {
 } from './mapping.js'
 import {
   type EncodingParams,
-  type GetMergeableRanksAsyncFn,
   type GetMergeableRanksFn,
   getEncodingParams,
-  getModelParamsAsync,
 } from './modelParams.js'
 import {
   EndOfPrompt,
@@ -129,6 +127,7 @@ export class GptEncoding {
     this.encodeChat = this.encodeChat.bind(this)
     this.encodeChatGenerator = this.encodeChatGenerator.bind(this)
     this.countTokens = this.countTokens.bind(this)
+    this.setMergeCacheSize = this.setMergeCacheSize.bind(this)
     this.modelName = modelName
   }
 
@@ -149,29 +148,6 @@ export class GptEncoding {
     return new GptEncoding({ ...modelParams, modelName })
   }
 
-  static async getEncodingApiAsync(
-    encodingName: EncodingName,
-    getMergeableRanks: GetMergeableRanksAsyncFn,
-  ): Promise<GptEncoding> {
-    const modelParams = await getModelParamsAsync(
-      encodingName,
-      getMergeableRanks,
-    )
-    return new GptEncoding(modelParams)
-  }
-
-  static async getEncodingApiForModelAsync(
-    modelName: ModelName,
-    getMergeableRanks: GetMergeableRanksAsyncFn,
-  ): Promise<GptEncoding> {
-    const encodingName = modelToEncodingMap[modelName]
-    const modelParams = await getModelParamsAsync(
-      encodingName,
-      getMergeableRanks,
-    )
-    return new GptEncoding({ ...modelParams, modelName })
-  }
-
   private processSpecialTokens({
     allowedSpecial,
     disallowedSpecial,
@@ -210,7 +186,9 @@ export class GptEncoding {
         allowedSpecial.forEach((val) => disallowedSpecialSet.delete(val))
         // disallowed takes precedence over allowed
         disallowedSpecial.forEach((val) => allowedSpecial.delete(val))
-        regexPattern = getSpecialTokenRegex(disallowedSpecial)
+        if (disallowedSpecial.size > 0) {
+          regexPattern = getSpecialTokenRegex(disallowedSpecial)
+        }
       } else {
         regexPattern = this.allSpecialTokenRegex
       }
@@ -354,11 +332,29 @@ export class GptEncoding {
    * Counts the number of tokens in the input.
    * @returns {number} The number of tokens.
    */
-  countTokens(input: string | Iterable<ChatMessage>): number {
-    const tokenGenerator =
-      typeof input === 'string'
-        ? this.encodeGenerator(input)
-        : this.encodeChatGenerator(input)
+  countTokens(
+    input: string | Iterable<ChatMessage>,
+    encodeOptions?: EncodeOptions,
+  ): number {
+    if (typeof input === 'string') {
+      const specialTokenConfig = encodeOptions
+        ? this.processSpecialTokens(encodeOptions)
+        : this.defaultSpecialTokenConfig
+
+      if (specialTokenConfig.regexPattern) {
+        const match = input.match(specialTokenConfig.regexPattern)
+        if (match !== null) {
+          throw new Error(`Disallowed special token found: ${match[0]}`)
+        }
+      }
+
+      return this.bytePairEncodingCoreProcessor.countNative(
+        input,
+        specialTokenConfig.allowedSpecial,
+      )
+    }
+
+    const tokenGenerator = this.encodeChatGenerator(input)
     let count = 0
     for (const tokens of tokenGenerator) {
       count += tokens.length
@@ -366,6 +362,10 @@ export class GptEncoding {
     return count
   }
 
+  setMergeCacheSize(size: number): void {
+    this.bytePairEncodingCoreProcessor.setMergeCacheSize(size)
+  }
+
   decode(inputTokensToDecode: Iterable<number>): string {
     return this.bytePairEncodingCoreProcessor.decodeNative(inputTokensToDecode)
   }

package/src/constants.ts

@@ -1 +1,2 @@
 export const ALL_SPECIAL_TOKENS = 'all'
+export const DEFAULT_MERGE_CACHE_SIZE = 100_000

package/src/encoding/cl100k_base.ts

@@ -17,6 +17,7 @@ const {
   encodeChat,
   encodeChatGenerator,
   vocabularySize,
+  setMergeCacheSize,
 } = api
 export {
   countTokens,
@@ -28,6 +29,7 @@ export {
   encodeChatGenerator,
   encodeGenerator,
   isWithinTokenLimit,
+  setMergeCacheSize,
   vocabularySize,
 }
 // eslint-disable-next-line import/no-default-export

package/src/encoding/o200k_base.ts

@@ -17,6 +17,7 @@ const {
   encodeChat,
   encodeChatGenerator,
   vocabularySize,
+  setMergeCacheSize,
 } = api
 export {
   countTokens,
@@ -28,6 +29,7 @@ export {
   encodeChatGenerator,
   encodeGenerator,
   isWithinTokenLimit,
+  setMergeCacheSize,
   vocabularySize,
 }
 // eslint-disable-next-line import/no-default-export

package/src/encoding/p50k_base.ts

@@ -15,6 +15,7 @@ const {
   isWithinTokenLimit,
   countTokens,
   vocabularySize,
+  setMergeCacheSize,
 } = api
 export {
   countTokens,
@@ -24,6 +25,7 @@ export {
   encode,
   encodeGenerator,
   isWithinTokenLimit,
+  setMergeCacheSize,
   vocabularySize,
 }
 // eslint-disable-next-line import/no-default-export

package/src/encoding/p50k_edit.ts

@@ -15,6 +15,7 @@ const {
   isWithinTokenLimit,
   countTokens,
   vocabularySize,
+  setMergeCacheSize,
 } = api
 export {
   countTokens,
@@ -24,6 +25,7 @@ export {
   encode,
   encodeGenerator,
   isWithinTokenLimit,
+  setMergeCacheSize,
   vocabularySize,
 }
 // eslint-disable-next-line import/no-default-export

package/src/encoding/r50k_base.ts

@@ -15,6 +15,7 @@ const {
   isWithinTokenLimit,
   countTokens,
   vocabularySize,
+  setMergeCacheSize,
 } = api
 export {
   countTokens,
@@ -24,6 +25,7 @@ export {
   encode,
   encodeGenerator,
   isWithinTokenLimit,
+  setMergeCacheSize,
   vocabularySize,
 }
 // eslint-disable-next-line import/no-default-export

package/src/encodingParams/{Cl100KBase.ts → cl100k_base.ts}

@@ -11,6 +11,7 @@ import {
   ImSep,
   ImStart,
 } from '../specialTokens.js'
+import { CL_AND_O_TOKEN_SPLIT_PATTERN } from './constants.js'
 
 export function Cl100KBase(
   bytePairRankDecoder: RawBytePairRanks,
@@ -27,8 +28,7 @@ export function Cl100KBase(
   ])
 
   return {
-    tokenSplitRegex:
-      /(?:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+/giu,
+    tokenSplitRegex: CL_AND_O_TOKEN_SPLIT_PATTERN,
     bytePairRankDecoder,
     specialTokensEncoder: specialTokenMapping,
   }

package/src/encodingParams/constants.ts

@@ -0,0 +1,5 @@
+export const R50K_TOKEN_SPLIT_REGEX =
+  /'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+/gu
+
+export const CL_AND_O_TOKEN_SPLIT_PATTERN =
+  /(?:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+/giu

package/src/encodingParams/{O200KBase.ts → o200k_base.ts}

@@ -11,6 +11,7 @@ import {
   ImSep,
   ImStart,
 } from '../specialTokens.js'
+import { CL_AND_O_TOKEN_SPLIT_PATTERN } from './constants.js'
 
 export function O200KBase(
   bytePairRankDecoder: RawBytePairRanks,
@@ -27,8 +28,7 @@ export function O200KBase(
   ])
 
   return {
-    tokenSplitRegex:
-      /(?:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+/giu,
+    tokenSplitRegex: CL_AND_O_TOKEN_SPLIT_PATTERN,
     bytePairRankDecoder,
     specialTokensEncoder: specialTokenMapping,
   }

package/src/encodingParams/{P50KBase.ts → p50k_base.ts}

@@ -1,14 +1,15 @@
 /* eslint-disable no-magic-numbers */
 import type { RawBytePairRanks } from '../BytePairEncodingCore.js'
-import { type EncodingParams, tokenSplitRegex } from '../modelParams.js'
+import { type EncodingParams } from '../modelParams.js'
 import { EndOfText } from '../specialTokens.js'
+import { R50K_TOKEN_SPLIT_REGEX } from './constants.js'
 
 export function P50KBase(
   bytePairRankDecoder: RawBytePairRanks,
 ): EncodingParams {
   return {
     expectedVocabularySize: 50_281,
-    tokenSplitRegex,
+    tokenSplitRegex: R50K_TOKEN_SPLIT_REGEX,
     bytePairRankDecoder,
     specialTokensEncoder: new Map<string, number>([[EndOfText, 50_256]]),
   }

package/src/encodingParams/{P50KEdit.ts → p50k_edit.ts}

@@ -1,7 +1,8 @@
 /* eslint-disable no-magic-numbers */
 import type { RawBytePairRanks } from '../BytePairEncodingCore.js'
-import { type EncodingParams, tokenSplitRegex } from '../modelParams.js'
+import { type EncodingParams } from '../modelParams.js'
 import { EndOfText, FimMiddle, FimPrefix, FimSuffix } from '../specialTokens.js'
+import { R50K_TOKEN_SPLIT_REGEX } from './constants.js'
 
 export function P50KEdit(
   bytePairRankDecoder: RawBytePairRanks,
@@ -14,7 +15,7 @@ export function P50KEdit(
   ])
 
   return {
-    tokenSplitRegex,
+    tokenSplitRegex: R50K_TOKEN_SPLIT_REGEX,
     bytePairRankDecoder,
     specialTokensEncoder: specialTokenMapping,
   }

package/src/encodingParams/{R50KBase.ts → r50k_base.ts}

@@ -1,14 +1,15 @@
 /* eslint-disable no-magic-numbers */
 import type { RawBytePairRanks } from '../BytePairEncodingCore.js'
-import { type EncodingParams, tokenSplitRegex } from '../modelParams.js'
+import { type EncodingParams } from '../modelParams.js'
 import { EndOfText } from '../specialTokens.js'
+import { R50K_TOKEN_SPLIT_REGEX } from './constants.js'
 
 export function R50KBase(
   bytePairRankDecoder: RawBytePairRanks,
 ): EncodingParams {
   return {
     expectedVocabularySize: 50_257,
-    tokenSplitRegex,
+    tokenSplitRegex: R50K_TOKEN_SPLIT_REGEX,
     bytePairRankDecoder,
     specialTokensEncoder: new Map<string, number>([[EndOfText, 50_256]]),
   }

package/src/extraApis.test.ts

@@ -0,0 +1,74 @@
+import { type ChatMessage, GptEncoding } from './GptEncoding.js'
+import { resolveEncoding } from './resolveEncoding.js'
+
+const sampleText = 'This is a test message.'
+const sampleChat: ChatMessage[] = [
+  {
+    role: 'system',
+    content: 'You are a helpful assistant.',
+  },
+  {
+    role: 'user',
+    content: 'Hello, how are you?',
+  },
+  {
+    role: 'assistant',
+    content: 'I am doing well, thank you for asking.',
+  },
+]
+
+describe('countTokens', () => {
+  const encoding = GptEncoding.getEncodingApiForModel(
+    'gpt-3.5-turbo',
+    resolveEncoding,
+  )
+
+  describe('text input', () => {
+    it('counts tokens in empty string', () => {
+      expect(encoding.countTokens('')).toBe(0)
+    })
+
+    it('counts tokens in simple text', () => {
+      expect(encoding.countTokens(sampleText)).toBe(
+        encoding.encode(sampleText).length,
+      )
+    })
+
+    it('counts tokens in text with special characters', () => {
+      const textWithSpecial = 'Hello 👋 world! 🌍'
+      expect(encoding.countTokens(textWithSpecial)).toBe(
+        encoding.encode(textWithSpecial).length,
+      )
+    })
+  })
+
+  describe('chat input', () => {
+    it('counts tokens in empty chat', () => {
+      expect(encoding.countTokens([])).toBe(3) // Due to assistant prompt tokens
+    })
+
+    it('counts tokens in sample chat', () => {
+      expect(encoding.countTokens(sampleChat)).toBe(
+        encoding.encodeChat(sampleChat).length,
+      )
+    })
+
+    it('matches token counts from encode methods', () => {
+      const tokens = encoding.encodeChat(sampleChat)
+      const count = encoding.countTokens(sampleChat)
+      expect(count).toBe(tokens.length)
+    })
+
+    it('counts tokens in single message chat', () => {
+      const singleMessage = [
+        {
+          role: 'user',
+          content: 'Hello world',
+        },
+      ] as const
+      expect(encoding.countTokens(singleMessage)).toBe(
+        encoding.encodeChat(singleMessage).length,
+      )
+    })
+  })
+})

package/src/model/gpt-3.5-turbo-0125.ts

@@ -13,11 +13,14 @@ const {
   encode,
   encodeGenerator,
   isWithinTokenLimit,
+  countTokens,
   encodeChat,
   encodeChatGenerator,
   vocabularySize,
+  setMergeCacheSize,
 } = api
 export {
+  countTokens,
   decode,
   decodeAsyncGenerator,
   decodeGenerator,
@@ -26,6 +29,7 @@ export {
   encodeChatGenerator,
   encodeGenerator,
   isWithinTokenLimit,
+  setMergeCacheSize,
   vocabularySize,
 }
 // eslint-disable-next-line import/no-default-export

package/src/model/gpt-3.5-turbo-0301.ts

@@ -13,11 +13,14 @@ const {
   encode,
   encodeGenerator,
   isWithinTokenLimit,
+  countTokens,
   encodeChat,
   encodeChatGenerator,
   vocabularySize,
+  setMergeCacheSize,
 } = api
 export {
+  countTokens,
   decode,
   decodeAsyncGenerator,
   decodeGenerator,
@@ -26,6 +29,7 @@ export {
   encodeChatGenerator,
   encodeGenerator,
   isWithinTokenLimit,
+  setMergeCacheSize,
   vocabularySize,
 }
 // eslint-disable-next-line import/no-default-export

package/src/model/gpt-3.5-turbo-0613.ts

@@ -13,11 +13,14 @@ const {
   encode,
   encodeGenerator,
   isWithinTokenLimit,
+  countTokens,
   encodeChat,
   encodeChatGenerator,
   vocabularySize,
+  setMergeCacheSize,
 } = api
 export {
+  countTokens,
   decode,
   decodeAsyncGenerator,
   decodeGenerator,
@@ -26,6 +29,7 @@ export {
   encodeChatGenerator,
   encodeGenerator,
   isWithinTokenLimit,
+  setMergeCacheSize,
   vocabularySize,
 }
 // eslint-disable-next-line import/no-default-export

package/src/model/gpt-3.5-turbo-1106.ts

@@ -13,11 +13,14 @@ const {
   encode,
   encodeGenerator,
   isWithinTokenLimit,
+  countTokens,
   encodeChat,
   encodeChatGenerator,
   vocabularySize,
+  setMergeCacheSize,
 } = api
 export {
+  countTokens,
   decode,
   decodeAsyncGenerator,
   decodeGenerator,
@@ -26,6 +29,7 @@ export {
   encodeChatGenerator,
   encodeGenerator,
   isWithinTokenLimit,
+  setMergeCacheSize,
   vocabularySize,
 }
 // eslint-disable-next-line import/no-default-export

package/src/model/gpt-3.5-turbo-16k-0613.ts

@@ -13,11 +13,14 @@ const {
   encode,
   encodeGenerator,
   isWithinTokenLimit,
+  countTokens,
   encodeChat,
   encodeChatGenerator,
   vocabularySize,
+  setMergeCacheSize,
 } = api
 export {
+  countTokens,
   decode,
   decodeAsyncGenerator,
   decodeGenerator,
@@ -26,6 +29,7 @@ export {
   encodeChatGenerator,
   encodeGenerator,
   isWithinTokenLimit,
+  setMergeCacheSize,
   vocabularySize,
 }
 // eslint-disable-next-line import/no-default-export