tiktoken-go

简体中文

OpenAI 的 tiktoken 的 Go 语言实现。

Tiktoken 是一个用于 OpenAI 模型的快速 BPE 分词器。

这是对原始 tiktoken 的移植。

使用方法

安装

go get github.com/pkoukk/tiktoken-go

缓存

Tiktoken-go 具有与原版 Tiktoken 库相同的缓存机制。

你可以通过设置环境变量 TIKTOKEN_CACHE_DIR 来指定缓存目录。

一旦设置了该变量，tiktoken-go 将使用此目录来缓存分词字典。

如果你没有设置这个环境变量，tiktoken-go 会在你首次初始化编码时每次都下载字典。

替代的 BPE 加载器

如果你不想使用缓存或每次都要下载字典，可以使用替代的 BPE 加载器。

只需在调用 tiktoken.GetEncoding 或 tiktoken.EncodingForModel 之前调用 tiktoken.SetBpeLoader 即可。

BpeLoader 是一个接口，你可以通过实现该接口来编写自己的 BPE 加载器。

离线 BPE 加载器

离线 BPE 加载器从嵌入文件中加载 BPE 字典，这在你不希望在运行时下载字典的情况下非常有用。

由于 BPE 字典的大小，这个加载器位于另一个项目中。

如果你需要这个加载器，请包含：tiktoken_loader

示例

通过编码获取 token

package main

import (
	"fmt"
	"github.com/pkoukk/tiktoken-go"
)

func main() {
	text := "Hello, world!"
	encoding := "cl100k_base"

	// 如果你不想在运行时下载字典，可以使用离线加载器
	// tiktoken.SetBpeLoader(tiktoken_loader.NewOfflineLoader())
	tke, err := tiktoken.GetEncoding(encoding)
	if err != nil {
		err = fmt.Errorf("getEncoding: %v", err)
		return
	}

	// 编码
	token := tke.Encode(text, nil, nil)

	// tokens
	fmt.Println(token)
	// num_tokens
	fmt.Println(len(token))
}

通过模型获取 token

package main

import (
	"fmt"
	"github.com/pkoukk/tiktoken-go"
)

func main() {
	text := "Hello, world!"
	encoding := "gpt-3.5-turbo"

	tkm, err := tiktoken.EncodingForModel(encoding)
	if err != nil {
		err = fmt.Errorf("getEncoding: %v", err)
		return
	}

	// 编码
	token := tkm.Encode(text, nil, nil)

	// tokens
	fmt.Println(token)
	// num_tokens
	fmt.Println(len(token))
}

计算聊天 API 调用的 token 数量

以下是用于计算传递给 gpt-3.5-turbo 或 gpt-4 的消息 token 数量的示例函数。

以下代码是基于 openai-cookbook 中的示例编写的，时间是 2023 年 6 月 28 日星期三。

请注意，消息的 token 计算方法可能会随时更改，因此这段代码在未来可能不再适用。

如果你需要精确的计算，请参考官方文档。

如果你发现这段代码不再适用，请随时提交 PR 或 Issue。

package main

import (
	"fmt"

	"github.com/pkoukk/tiktoken-go"
	"github.com/sashabaranov/go-openai"
)

// OpenAI Cookbook: https://github.com/openai/openai-cookbook/blob/main/examples/How_to_count_tokens_with_tiktoken.ipynb
func NumTokensFromMessages(messages []openai.ChatCompletionMessage, model string) (numTokens int) {
	tkm, err := tiktoken.EncodingForModel(model)
	if err != nil {
		err = fmt.Errorf("encoding for model: %v", err)
		log.Println(err)
		return
	}

	var tokensPerMessage, tokensPerName int
	switch model {
	case "gpt-3.5-turbo-0613",
		"gpt-3.5-turbo-16k-0613",
		"gpt-4-0314",
		"gpt-4-32k-0314",
		"gpt-4-0613",
		"gpt-4-32k-0613":
		tokensPerMessage = 3
		tokensPerName = 1
	case "gpt-3.5-turbo-0301":
		tokensPerMessage = 4 // 每条消息遵循 <|start|>{role/name}\n{content}<|end|>\n
		tokensPerName = -1   // 如果有 name，则省略 role
	default:
		if strings.Contains(model, "gpt-3.5-turbo") {
			log.Println("警告：gpt-3.5-turbo 可能会随时间更新。返回假设为 gpt-3.5-turbo-0613 的 token 数量。")
			return NumTokensFromMessages(messages, "gpt-3.5-turbo-0613")
		} else if strings.Contains(model, "gpt-4") {
			log.Println("警告：gpt-4 可能会随时间更新。返回假设为 gpt-4-0613 的 token 数量。")
			return NumTokensFromMessages(messages, "gpt-4-0613")
		} else {
			err = fmt.Errorf("num_tokens_from_messages() 对于模型 %s 尚未实现。请参阅 https://github.com/openai/openai-python/blob/main/chatml.md，了解消息如何转换为 token 的信息。")
			log.Println(err)
			return
		}
	}

	for _, message := range messages {
		numTokens += tokensPerMessage
		numTokens += len(tkm.Encode(message.Content, nil, nil))
		numTokens += len(tkm.Encode(message.Role, nil, nil))
		numTokens += len(tkm.Encode(message.Name, nil, nil))
		if message.Name != "" {
			numTokens += tokensPerName
		}
	}
	numTokens += 3 // 每个回复都以 <|start|>assistant<|message|> 开头
	return numTokens
}

可用的编码

编码名称	OpenAI 模型
o200k_base	gpt-4o, gpt-4.1, gpt-4.5
cl100k_base	gpt-4, gpt-3.5-turbo, text-embedding-ada-002, text-embedding-3-small, text-embedding-3-large
p50k_base	Codex 模型、text-davinci-002、text-davinci-003
r50k_base (或 gpt2)	GPT-3 模型，如 davinci

可用模型

模型名称	OpenAI 模型
gpt-4.5-*	o200k_base
gpt-4.1-*	o200k_base
gpt-4o-*	o200k_base
gpt-4-*	cl100k_base
gpt-3.5-turbo-*	cl100k_base
gpt-4o	o200k_base
gpt-4	cl100k_base
gpt-3.5-turbo	cl100k_base
text-davinci-003	p50k_base
text-davinci-002	p50k_base
text-davinci-001	r50k_base
text-curie-001	r50k_base
text-babbage-001	r50k_base
text-ada-001	r50k_base
davinci	r50k_base
curie	r50k_base
babbage	r50k_base
ada	r50k_base
code-davinci-002	p50k_base
code-davinci-001	p50k_base
code-cushman-002	p50k_base
code-cushman-001	p50k_base
davinci-codex	p50k_base
cushman-codex	p50k_base
text-davinci-edit-001	p50k_edit
code-davinci-edit-001	p50k_edit
text-embedding-ada-002	cl100k_base
text-embedding-3-small	cl100k_base
text-embedding-3-large	cl100k_base
text-similarity-davinci-001	r50k_base
text-similarity-curie-001	r50k_base
text-similarity-babbage-001	r50k_base
text-similarity-ada-001	r50k_base
text-search-davinci-doc-001	r50k_base
text-search-curie-doc-001	r50k_base
text-search-babbage-doc-001	r50k_base
text-search-ada-doc-001	r50k_base
code-search-babbage-code-001	r50k_base
code-search-ada-code-001	r50k_base
gpt2	gpt2

测试

你可以在 test 文件夹中运行测试

与原始 tiktoken 对比

通过编码获取 token

结果

通过模型获取 token

结果

基准测试

你可以在 test 文件夹中运行基准测试

基准测试结果

名称	time/op	操作系统	CPU	文本	次数
tiktoken-go	8795ns	macOS 13.2	Apple M1	UDHR	100000
tiktoken	8838ns	macOS 13.2	Apple M1	UDHR	100000

看起来性能几乎相同。

也许差异是由于机器性能的不同造成的。

或者我的基准测试方法可能不太合适。

如果你有更好的基准测试方法，或者想添加你的基准测试结果，请随时提交 PR。

对于新的 o200k_base 编码，它似乎比 cl100k_base 更慢。在以下基准测试中，tiktoken-go 略慢于 tiktoken。

名称	编码	time/op	操作系统	CPU	文本	次数
tiktoken-go	o200k_base	108522 ns	Ubuntu 22.04	AMD Ryzen 9 5900HS	UDHR	100000
tiktoken	o200k_base	70198 ns	Ubuntu 22.04	AMD Ryzen 9 5900HS	UDHR	100000
tiktoken-go	cl100k_base	94502 ns	Ubuntu 22.04	AMD Ryzen 9 5900HS	UDHR	100000
tiktoken	cl100k_base	54642 ns	Ubuntu 22.04	AMD Ryzen 9 5900HS	UDHR	100000

许可证

MIT

tiktoken-go 快速上手指南

tiktoken-go 是 OpenAI 官方 tiktoken 库的 Go 语言移植版本，用于快速计算 OpenAI 模型（如 GPT-4、GPT-3.5-Turbo）的 Token 数量。

环境准备

• 操作系统：Linux, macOS, Windows
• Go 版本：建议 Go 1.18 及以上版本
• 前置依赖：无特殊系统依赖，仅需标准的 Go 开发环境

安装步骤

使用 Go 模块直接安装：

go get github.com/pkoukk/tiktoken-go

提示：如果下载速度较慢，可配置国内 Go 代理：

export GOPROXY=https://goproxy.cn,direct
go get github.com/pkoukk/tiktoken-go

基本使用

1. 通过编码名称获取 Token

最基础的用法是指定编码名称（如 cl100k_base）对文本进行编码。

package main

import (
	"fmt"
	"github.com/pkoukk/tiktoken-go"
)

func main() {
	text := "Hello, world!"
	encodingName := "cl100k_base"

	// 获取编码器
	tke, err := tiktoken.GetEncoding(encodingName)
	if err != nil {
		fmt.Printf("getEncoding: %v\n", err)
		return
	}

	// 编码文本
	token := tke.Encode(text, nil, nil)

	// 输出结果
	fmt.Println(token)        // Token ID 列表
	fmt.Println(len(token))   // Token 数量
}

2. 通过模型名称获取 Token

直接指定 OpenAI 模型名称（如 gpt-3.5-turbo），库会自动匹配对应的编码规则。

package main

import (
	"fmt"
	"github.com/pkoukk/tiktoken-go"
)

func main() {
	text := "Hello, world!"
	modelName := "gpt-3.5-turbo"

	// 根据模型获取编码器
	tkm, err := tiktoken.EncodingForModel(modelName)
	if err != nil {
		fmt.Printf("getEncoding: %v\n", err)
		return
	}

	// 编码文本
	token := tkm.Encode(text, nil, nil)

	// 输出结果
	fmt.Println(token)        // Token ID 列表
	fmt.Println(len(token))   // Token 数量
}

注意事项

• 缓存机制：首次运行时会自动下载编码字典。可通过设置环境变量 TIKTOKEN_CACHE_DIR 指定缓存目录，避免重复下载。
• 离线使用：若需完全离线运行，可结合 tiktoken-go-loader 项目加载本地字典文件。