From 9491b753c34ac90e634d8adda141e8bf9051be9b Mon Sep 17 00:00:00 2001 From: Danny Avila <110412045+danny-avila@users.noreply.github.com> Date: Thu, 14 Sep 2023 19:40:21 -0400 Subject: [PATCH] =?UTF-8?q?fix:=20Match=20OpenAI=20Token=20Counting=20Stra?= =?UTF-8?q?tegy=20=F0=9F=AA=99=20(#945)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * wip token fix * fix: complete token count refactor to match OpenAI example * chore: add back sendPayload method (accidentally deleted) * chore: revise JSDoc for getTokenCountForMessage --- api/app/clients/AnthropicClient.js | 5 +- api/app/clients/BaseClient.js | 47 +++++++---------- api/app/clients/ChatGPTClient.js | 45 +++++++---------- api/app/clients/GoogleClient.js | 5 +- api/app/clients/OpenAIClient.js | 5 +- api/app/clients/specs/BaseClient.test.js | 7 ++- api/app/clients/specs/OpenAIClient.test.js | 59 ++++++++++++++++++++++ api/package.json | 2 +- api/server/routes/tokenizer.js | 8 +-- api/utils/tokens.js | 1 + package-lock.json | 7 ++- 11 files changed, 115 insertions(+), 76 deletions(-) diff --git a/api/app/clients/AnthropicClient.js b/api/app/clients/AnthropicClient.js index 5c35bc28f4..4c871afb9d 100644 --- a/api/app/clients/AnthropicClient.js +++ b/api/app/clients/AnthropicClient.js @@ -1,9 +1,6 @@ // const { Agent, ProxyAgent } = require('undici'); const BaseClient = require('./BaseClient'); -const { - encoding_for_model: encodingForModel, - get_encoding: getEncoding, -} = require('@dqbd/tiktoken'); +const { encoding_for_model: encodingForModel, get_encoding: getEncoding } = require('tiktoken'); const Anthropic = require('@anthropic-ai/sdk'); const HUMAN_PROMPT = '\n\nHuman:'; diff --git a/api/app/clients/BaseClient.js b/api/app/clients/BaseClient.js index bfb65cf93c..e90d9fefb0 100644 --- a/api/app/clients/BaseClient.js +++ b/api/app/clients/BaseClient.js @@ -272,7 +272,9 @@ class BaseClient { * @returns {Object} An object with three properties: `context`, `remainingContextTokens`, and `messagesToRefine`. `context` is an array of messages that fit within the token limit. `remainingContextTokens` is the number of tokens remaining within the limit after adding the messages to the context. `messagesToRefine` is an array of messages that were not added to the context because they would have exceeded the token limit. */ async getMessagesWithinTokenLimit(messages) { - let currentTokenCount = 0; + // Every reply is primed with <|start|>assistant<|message|>, so we + // start with 3 tokens for the label after all messages have been counted. + let currentTokenCount = 3; let context = []; let messagesToRefine = []; let refineIndex = -1; @@ -562,44 +564,29 @@ class BaseClient { * Algorithm adapted from "6. Counting tokens for chat API calls" of * https://github.com/openai/openai-cookbook/blob/main/examples/How_to_count_tokens_with_tiktoken.ipynb * - * An additional 2 tokens need to be added for metadata after all messages have been counted. + * An additional 3 tokens need to be added for assistant label priming after all messages have been counted. * - * @param {*} message + * @param {Object} message */ getTokenCountForMessage(message) { - let tokensPerMessage; - let nameAdjustment; - if (this.modelOptions.model.startsWith('gpt-4')) { - tokensPerMessage = 3; - nameAdjustment = 1; - } else { + // Note: gpt-3.5-turbo and gpt-4 may update over time. Use default for these as well as for unknown models + let tokensPerMessage = 3; + let tokensPerName = 1; + + if (this.modelOptions.model === 'gpt-3.5-turbo-0301') { tokensPerMessage = 4; - nameAdjustment = -1; + tokensPerName = -1; } - if (this.options.debug) { - console.debug('getTokenCountForMessage', message); - } - - // Map each property of the message to the number of tokens it contains - const propertyTokenCounts = Object.entries(message).map(([key, value]) => { - if (key === 'tokenCount' || typeof value !== 'string') { - return 0; + let numTokens = tokensPerMessage; + for (let [key, value] of Object.entries(message)) { + numTokens += this.getTokenCount(value); + if (key === 'name') { + numTokens += tokensPerName; } - // Count the number of tokens in the property value - const numTokens = this.getTokenCount(value); - - // Adjust by `nameAdjustment` tokens if the property key is 'name' - const adjustment = key === 'name' ? nameAdjustment : 0; - return numTokens + adjustment; - }); - - if (this.options.debug) { - console.debug('propertyTokenCounts', propertyTokenCounts); } - // Sum the number of tokens in all properties and add `tokensPerMessage` for metadata - return propertyTokenCounts.reduce((a, b) => a + b, tokensPerMessage); + return numTokens; } async sendPayload(payload, opts = {}) { diff --git a/api/app/clients/ChatGPTClient.js b/api/app/clients/ChatGPTClient.js index 72715669e6..5fa6344d2c 100644 --- a/api/app/clients/ChatGPTClient.js +++ b/api/app/clients/ChatGPTClient.js @@ -1,9 +1,6 @@ const crypto = require('crypto'); const Keyv = require('keyv'); -const { - encoding_for_model: encodingForModel, - get_encoding: getEncoding, -} = require('@dqbd/tiktoken'); +const { encoding_for_model: encodingForModel, get_encoding: getEncoding } = require('tiktoken'); const { fetchEventSource } = require('@waylaidwanderer/fetch-event-source'); const { Agent, ProxyAgent } = require('undici'); const BaseClient = require('./BaseClient'); @@ -526,8 +523,8 @@ ${botMessage.message} const prompt = `${promptBody}${promptSuffix}`; if (isChatGptModel) { messagePayload.content = prompt; - // Add 2 tokens for metadata after all messages have been counted. - currentTokenCount += 2; + // Add 3 tokens for Assistant Label priming after all messages have been counted. + currentTokenCount += 3; } // Use up to `this.maxContextTokens` tokens (prompt + response), but try to leave `this.maxTokens` tokens for the response. @@ -554,33 +551,29 @@ ${botMessage.message} * Algorithm adapted from "6. Counting tokens for chat API calls" of * https://github.com/openai/openai-cookbook/blob/main/examples/How_to_count_tokens_with_tiktoken.ipynb * - * An additional 2 tokens need to be added for metadata after all messages have been counted. + * An additional 3 tokens need to be added for assistant label priming after all messages have been counted. * - * @param {*} message + * @param {Object} message */ getTokenCountForMessage(message) { - let tokensPerMessage; - let nameAdjustment; - if (this.modelOptions.model.startsWith('gpt-4')) { - tokensPerMessage = 3; - nameAdjustment = 1; - } else { + // Note: gpt-3.5-turbo and gpt-4 may update over time. Use default for these as well as for unknown models + let tokensPerMessage = 3; + let tokensPerName = 1; + + if (this.modelOptions.model === 'gpt-3.5-turbo-0301') { tokensPerMessage = 4; - nameAdjustment = -1; + tokensPerName = -1; } - // Map each property of the message to the number of tokens it contains - const propertyTokenCounts = Object.entries(message).map(([key, value]) => { - // Count the number of tokens in the property value - const numTokens = this.getTokenCount(value); + let numTokens = tokensPerMessage; + for (let [key, value] of Object.entries(message)) { + numTokens += this.getTokenCount(value); + if (key === 'name') { + numTokens += tokensPerName; + } + } - // Adjust by `nameAdjustment` tokens if the property key is 'name' - const adjustment = key === 'name' ? nameAdjustment : 0; - return numTokens + adjustment; - }); - - // Sum the number of tokens in all properties and add `tokensPerMessage` for metadata - return propertyTokenCounts.reduce((a, b) => a + b, tokensPerMessage); + return numTokens; } } diff --git a/api/app/clients/GoogleClient.js b/api/app/clients/GoogleClient.js index dee0ab829d..d3f77d0690 100644 --- a/api/app/clients/GoogleClient.js +++ b/api/app/clients/GoogleClient.js @@ -1,10 +1,7 @@ const BaseClient = require('./BaseClient'); const { google } = require('googleapis'); const { Agent, ProxyAgent } = require('undici'); -const { - encoding_for_model: encodingForModel, - get_encoding: getEncoding, -} = require('@dqbd/tiktoken'); +const { encoding_for_model: encodingForModel, get_encoding: getEncoding } = require('tiktoken'); const tokenizersCache = {}; diff --git a/api/app/clients/OpenAIClient.js b/api/app/clients/OpenAIClient.js index e3d15e66e2..0fbb0339e5 100644 --- a/api/app/clients/OpenAIClient.js +++ b/api/app/clients/OpenAIClient.js @@ -1,9 +1,6 @@ const BaseClient = require('./BaseClient'); const ChatGPTClient = require('./ChatGPTClient'); -const { - encoding_for_model: encodingForModel, - get_encoding: getEncoding, -} = require('@dqbd/tiktoken'); +const { encoding_for_model: encodingForModel, get_encoding: getEncoding } = require('tiktoken'); const { maxTokensMap, genAzureChatCompletion } = require('../../utils'); const { runTitleChain } = require('./chains'); const { createLLM } = require('./llm'); diff --git a/api/app/clients/specs/BaseClient.test.js b/api/app/clients/specs/BaseClient.test.js index 10d5868cb0..183aaae989 100644 --- a/api/app/clients/specs/BaseClient.test.js +++ b/api/app/clients/specs/BaseClient.test.js @@ -138,7 +138,8 @@ describe('BaseClient', () => { { role: 'assistant', content: 'How can I help you?', tokenCount: 19 }, { role: 'user', content: 'I have a question.', tokenCount: 18 }, ]; - const expectedRemainingContextTokens = 58; // 100 - 5 - 19 - 18 + // Subtract 3 tokens for Assistant Label priming after all messages have been counted. + const expectedRemainingContextTokens = 58 - 3; // (100 - 5 - 19 - 18) - 3 const expectedMessagesToRefine = []; const result = await TestClient.getMessagesWithinTokenLimit(messages); @@ -168,7 +169,9 @@ describe('BaseClient', () => { { role: 'assistant', content: 'How can I help you?', tokenCount: 19 }, { role: 'user', content: 'I have a question.', tokenCount: 18 }, ]; - const expectedRemainingContextTokens = 8; // 50 - 18 - 19 - 5 + + // Subtract 3 tokens for Assistant Label priming after all messages have been counted. + const expectedRemainingContextTokens = 8 - 3; // (50 - 18 - 19 - 5) - 3 const expectedMessagesToRefine = [ { role: 'user', content: 'I need a coffee, stat!', tokenCount: 30 }, { role: 'assistant', content: 'Sure, I can help with that.', tokenCount: 30 }, diff --git a/api/app/clients/specs/OpenAIClient.test.js b/api/app/clients/specs/OpenAIClient.test.js index dd4de5cc7c..06874b2e17 100644 --- a/api/app/clients/specs/OpenAIClient.test.js +++ b/api/app/clients/specs/OpenAIClient.test.js @@ -213,4 +213,63 @@ describe('OpenAIClient', () => { expect(result.prompt).toEqual([]); }); }); + + describe('getTokenCountForMessage', () => { + const example_messages = [ + { + role: 'system', + content: + 'You are a helpful, pattern-following assistant that translates corporate jargon into plain English.', + }, + { + role: 'system', + name: 'example_user', + content: 'New synergies will help drive top-line growth.', + }, + { + role: 'system', + name: 'example_assistant', + content: 'Things working well together will increase revenue.', + }, + { + role: 'system', + name: 'example_user', + content: + 'Let\'s circle back when we have more bandwidth to touch base on opportunities for increased leverage.', + }, + { + role: 'system', + name: 'example_assistant', + content: 'Let\'s talk later when we\'re less busy about how to do better.', + }, + { + role: 'user', + content: + 'This late pivot means we don\'t have time to boil the ocean for the client deliverable.', + }, + ]; + + const testCases = [ + { model: 'gpt-3.5-turbo-0301', expected: 127 }, + { model: 'gpt-3.5-turbo-0613', expected: 129 }, + { model: 'gpt-3.5-turbo', expected: 129 }, + { model: 'gpt-4-0314', expected: 129 }, + { model: 'gpt-4-0613', expected: 129 }, + { model: 'gpt-4', expected: 129 }, + { model: 'unknown', expected: 129 }, + ]; + + testCases.forEach((testCase) => { + it(`should return ${testCase.expected} tokens for model ${testCase.model}`, () => { + client.modelOptions.model = testCase.model; + client.selectTokenizer(); + // 3 tokens for assistant label + let totalTokens = 3; + for (let message of example_messages) { + totalTokens += client.getTokenCountForMessage(message); + } + expect(totalTokens).toBe(testCase.expected); + }); + }); + }); }); diff --git a/api/package.json b/api/package.json index 3ccb01cb54..5f8f1b94d8 100644 --- a/api/package.json +++ b/api/package.json @@ -23,7 +23,6 @@ "dependencies": { "@anthropic-ai/sdk": "^0.5.4", "@azure/search-documents": "^11.3.2", - "@dqbd/tiktoken": "^1.0.7", "@keyv/mongo": "^2.1.8", "@waylaidwanderer/chatgpt-api": "^1.37.2", "axios": "^1.3.4", @@ -60,6 +59,7 @@ "passport-local": "^1.0.0", "pino": "^8.12.1", "sharp": "^0.32.5", + "tiktoken": "^1.0.10", "ua-parser-js": "^1.0.36", "zod": "^3.22.2" }, diff --git a/api/server/routes/tokenizer.js b/api/server/routes/tokenizer.js index 9fd79d0acf..240faa94ee 100644 --- a/api/server/routes/tokenizer.js +++ b/api/server/routes/tokenizer.js @@ -1,9 +1,9 @@ const express = require('express'); const router = express.Router(); -const { Tiktoken } = require('@dqbd/tiktoken/lite'); -const { load } = require('@dqbd/tiktoken/load'); -const registry = require('@dqbd/tiktoken/registry.json'); -const models = require('@dqbd/tiktoken/model_to_encoding.json'); +const { Tiktoken } = require('tiktoken/lite'); +const { load } = require('tiktoken/load'); +const registry = require('tiktoken/registry.json'); +const models = require('tiktoken/model_to_encoding.json'); const requireJwtAuth = require('../middleware/requireJwtAuth'); router.post('/', requireJwtAuth, async (req, res) => { diff --git a/api/utils/tokens.js b/api/utils/tokens.js index 7d0cb02377..eebe56e5a9 100644 --- a/api/utils/tokens.js +++ b/api/utils/tokens.js @@ -41,6 +41,7 @@ const maxTokensMap = { 'gpt-4': 8191, 'gpt-4-0613': 8191, 'gpt-4-32k': 32767, + 'gpt-4-32k-0314': 32767, 'gpt-4-32k-0613': 32767, 'gpt-3.5-turbo': 4095, 'gpt-3.5-turbo-0613': 4095, diff --git a/package-lock.json b/package-lock.json index e2a1865b55..7abad398bb 100644 --- a/package-lock.json +++ b/package-lock.json @@ -44,7 +44,6 @@ "dependencies": { "@anthropic-ai/sdk": "^0.5.4", "@azure/search-documents": "^11.3.2", - "@dqbd/tiktoken": "^1.0.7", "@keyv/mongo": "^2.1.8", "@waylaidwanderer/chatgpt-api": "^1.37.2", "axios": "^1.3.4", @@ -81,6 +80,7 @@ "passport-local": "^1.0.0", "pino": "^8.12.1", "sharp": "^0.32.5", + "tiktoken": "^1.0.10", "ua-parser-js": "^1.0.36", "zod": "^3.22.2" }, @@ -21886,6 +21886,11 @@ "real-require": "^0.2.0" } }, + "node_modules/tiktoken": { + "version": "1.0.10", + "resolved": "https://registry.npmjs.org/tiktoken/-/tiktoken-1.0.10.tgz", + "integrity": "sha512-gF8ndTCNu7WcRFbl1UUWaFIB4CTXmHzS3tRYdyUYF7x3C6YR6Evoao4zhKDmWIwv2PzNbzoQMV8Pxt+17lEDbA==" + }, "node_modules/tmp": { "version": "0.0.33", "resolved": "https://registry.npmjs.org/tmp/-/tmp-0.0.33.tgz",