| 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 countTokens = async (text = '', modelName = 'gpt-3.5-turbo') => { | |
| let encoder = null; | |
| try { | |
| const model = await load(registry[models[modelName]]); | |
| encoder = new Tiktoken(model.bpe_ranks, model.special_tokens, model.pat_str); | |
| const tokens = encoder.encode(text); | |
| encoder.free(); | |
| return tokens.length; | |
| } catch (e) { | |
| console.error(e); | |
| if (encoder) { | |
| encoder.free(); | |
| } | |
| return 0; | |
| } | |
| }; | |
| module.exports = countTokens; | |