konnosDashboard
tokenizer.ts94 lines · main
1import { getEncoding } from 'js-tiktoken'
2import type OpenAI from 'openai'
3
4export const tokenizer = getEncoding('cl100k_base')
5
6/**
7 * Count the tokens for multi-message chat completion requests
8 */
9export function getChatRequestTokenCount(
10 messages: OpenAI.Chat.Completions.ChatCompletionMessageParam[],
11 model = 'gpt-4o-mini-2024-07-18'
12): number {
13 const tokensPerRequest = 3 // every reply is primed with <|im_start|>assistant<|im_sep|>
14 const numTokens = messages.reduce((acc, message) => acc + getMessageTokenCount(message, model), 0)
15
16 return numTokens + tokensPerRequest
17}
18
19/**
20 * Count the tokens for a single message within a chat completion request
21 *
22 * See "Counting tokens for chat API calls"
23 * from https://github.com/openai/openai-cookbook/blob/834181d5739740eb8380096dac7056c925578d9a/examples/How_to_count_tokens_with_tiktoken.ipynb
24 */
25export function getMessageTokenCount(
26 message: OpenAI.Chat.Completions.ChatCompletionMessageParam,
27 model = 'gpt-4o-mini-2024-07-18'
28): number {
29 let tokensPerMessage: number
30 let tokensPerName: number
31
32 switch (model) {
33 case 'gpt-3.5-turbo':
34 console.warn(
35 'Warning: gpt-3.5-turbo may change over time. Returning num tokens assuming gpt-3.5-turbo-0301.'
36 )
37 return getMessageTokenCount(message, 'gpt-3.5-turbo-0301')
38 case 'gpt-4':
39 console.warn('Warning: gpt-4 may change over time. Returning num tokens assuming gpt-4-0314.')
40 return getMessageTokenCount(message, 'gpt-4-0314')
41 case 'gpt-3.5-turbo-0301':
42 tokensPerMessage = 4 // every message follows <|start|>{role/name}\n{content}<|end|>\n
43 tokensPerName = -1 // if there's a name, the role is omitted
44 break
45 case 'gpt-4o-mini-2024-07-18':
46 tokensPerMessage = 3
47 tokensPerName = 1
48 break
49 case 'gpt-4-0314':
50 tokensPerMessage = 3
51 tokensPerName = 1
52 break
53 default:
54 throw new Error(
55 `Unknown model '${model}'. See https://github.com/openai/openai-python/blob/main/chatml.md for information on how messages are converted to tokens.`
56 )
57 }
58
59 return Object.entries(message).reduce((acc, [key, value]) => {
60 acc += tokenizer.encode(value).length
61 if (key === 'name') {
62 acc += tokensPerName
63 }
64 return acc
65 }, tokensPerMessage)
66}
67
68/**
69 * Get the maximum number of tokens for a model's context.
70 *
71 * Includes tokens in both message and completion.
72 */
73export function getMaxTokenCount(model: string): number {
74 switch (model) {
75 case 'gpt-3.5-turbo':
76 console.warn(
77 'Warning: gpt-3.5-turbo may change over time. Returning max num tokens assuming gpt-3.5-turbo-0301.'
78 )
79 return getMaxTokenCount('gpt-3.5-turbo-0301')
80 case 'gpt-4':
81 console.warn(
82 'Warning: gpt-4 may change over time. Returning max num tokens assuming gpt-4-0314.'
83 )
84 return getMaxTokenCount('gpt-4-0314')
85 case 'gpt-3.5-turbo-0301':
86 return 4097
87 case 'gpt-4-0314':
88 return 4097
89 case 'gpt-4o-mini-2024-07-18':
90 return 4097
91 default:
92 throw new Error(`Unknown model '${model}'`)
93 }
94}