Tiktoken：OpenAI 开源 Token 分词工具

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

import tiktoken
enc = tiktoken.get_encoding("o200k_base")
assert enc.decode(enc.encode("hello world")) == "hello world"

# To get the tokeniser corresponding to a specific model in the OpenAI API:
enc = tiktoken.encoding_for_model("gpt-4o")

tiktoken 的开源版本可以从 PyPI 安装：

pip install tiktoken

分词器 API 在 tiktoken/core.py 中有文档说明。

使用 tiktoken 的示例代码可以在 OpenAI Cookbook 中找到。

tiktoken 比同类开源分词器快 3-6 倍：

性能测试基于 1GB 文本，使用 GPT-2 分词器，采用 tokenizers==0.13.2、transformers==4.24.0 和 tiktoken==0.2.0。

请在问题追踪器中提出问题。

如果你在 OpenAI 工作，请务必查看内部文档或随时联系 @shantanu。

BPE 和语言模型分词

语言模型看待文本的方式与你我不同，它们看到的是一系列数字（称为 token）。字节对编码（BPE）是一种将文本转换为 token 的方法。它有几个理想的特性：

它是可逆且无损的，所以你可以将 token 转换回原始文本

它适用于任意文本，即使是分词器训练数据中没有的文本

它可以压缩文本：token 序列比原始文本对应的字节更短。实际上，平均而言，每个 token 对应约 4 个字节

它尝试让模型看到常见的子词。例如，"ing" 是英文中的常见子词，所以 BPE 编码通常会将 "encoding" 拆分成 "encod" 和 "ing" 这样的 token（而不是比如 "enc" 和 "oding"）。因为模型会在不同的上下文中反复看到 "ing" token，这帮助模型泛化并更好地理解语法。

tiktoken 包含一个教育子模块，如果你想更多地了解 BPE 的细节，它会更友好。它包括帮助可视化 BPE 过程的代码：

from tiktoken._educational import *

# Train a BPE tokeniser on a small amount of text
enc = train_simple_encoding()

# Visualise how the GPT-4 encoder encodes text
enc = SimpleBytePairEncoding.from_tiktoken("cl100k_base")
enc.encode("hello world aaaaaaaaaaaa")

扩展 tiktoken

你可能希望扩展 tiktoken 以支持新的编码。有两种方法可以做到这一点。

方法一：直接创建 Encoding 对象

创建你的 Encoding 对象，按照你想要的方式，简单地传递它。

cl100k_base = tiktoken.get_encoding("cl100k_base")

# In production, load the arguments directly instead of accessing private attributes
# See openai_public.py for examples of arguments for specific encodings
enc = tiktoken.Encoding(
    # If you're changing the set of special tokens, make sure to use a different name
    # It should be clear from the name what behaviour to expect.
    name="cl100k_im",
    pat_str=cl100k_base._pat_str,
    mergeable_ranks=cl100k_base._mergeable_ranks,
    special_tokens={
        **cl100k_base._special_tokens,
        "<|im_start|>": 100264,
        "<|im_end|>": 100265,
    }
)

方法二：使用 tiktoken_ext 插件机制

使用 tiktoken_ext 插件机制向 tiktoken 注册你的 Encoding 对象。

这仅在你需要 tiktoken.get_encoding 找到你的编码时才有用，否则首选方案 1。

要这样做，你需要在 tiktoken_ext 下创建一个命名空间包。

像这样布置你的项目，确保省略 tiktoken_ext/init.py 文件：

my_tiktoken_extension
├── tiktoken_ext
│   └── my_encodings.py
└── setup.py

my_encodings.py 应该是一个模块，其中包含一个名为 ENCODING_CONSTRUCTORS 的变量。这是一个从编码名称映射到函数的字典，该函数不接收任何参数，并返回可以传递给 tiktoken.Encoding 来构造该编码的参数。有关示例，请参阅 tiktoken_ext/openai_public.py。有关精确细节，请参阅 tiktoken/registry.py。

你的 setup.py 应该类似于这样：

from setuptools import setup, find_namespace_packages

setup(
    name="my_tiktoken_extension",
    packages=find_namespace_packages(include=['tiktoken_ext*']),
    install_requires=["tiktoken"],
    ...
)

然后简单地运行 pip install ./my_tiktoken_extension，你应该能够使用你的自定义编码！确保不要使用可编辑安装。