当前位置：物联沃-IOTWORD物联网 > 技术教程 > 【Python】tiktoken库：OpenAI开发的BPE分词器库详解与快速字节对编码实践
代码收藏家代码收藏家 技术教程

【Python】tiktoken库：OpenAI开发的BPE分词器库详解与快速字节对编码实践

tiktoken 是一个由 OpenAI 开发的快速字节对编码（Byte Pair Encoding, BPE）分词器库，专门为 OpenAI 的语言模型（如 GPT-3、GPT-4）设计。它用于将文本转换为模型可处理的 token（标记），并支持 token 计数、编码和解码操作。tiktoken 以高效性和速度著称，性能比其他开源分词器（如 Hugging Face 的 GPT2TokenizerFast）快 3-6 倍，广泛应用于自然语言处理（NLP）、API 成本估算和输入验证。

以下是对 tiktoken 库的详细介绍，包括其功能、用法、实际应用和注意事项。

1. tiktoken 库的作用

  • 文本分词：将文本字符串分割为 token，用于输入到语言模型。
  • token 计数：计算文本的 token 数量，帮助验证是否符合模型的 token 限制（如 GPT-3.5-turbo 的 16k token 限制）。
  • 成本估算：根据 token 数量估算 OpenAI API 调用成本（API 按 token 计费）。
  • 高效编码：支持多种编码方式（如 cl100k_base、o200k_base），适配不同模型。
  • 跨平台：不仅支持 Python，还提供 JavaScript、Dart 等语言的端口。
  • 教育功能：包含 _educational 子模块，帮助理解 BPE 分词原理。

    • 关键场景：

  • 验证输入文本是否超过模型的 token 限制。
  • 估算 API 调用成本。
  • 处理大批量文本分词任务。
  • 学习和调试 BPE 分词算法。

    • 2. 安装与环境要求

  • Python 版本：支持 Python 3.9+（推荐 3.10+）。

  • 依赖：

  • regex：用于正则表达式处理。
  • requests：下载编码文件（首次使用）。
  • 可选：blobfile（处理大文件）。

  • 安装命令：

    pip install tiktoken

  • 验证安装：

    import tiktoken
print(tiktoken.__version__)  # 示例输出: 0.9.0

  • Rust 编译器（若从源代码安装）：

  • 部分环境中需安装 Rust 编译器（通过 rustup 安装）。
  • 若遇到安装错误，可设置环境变量： 
    export PYO3_USE_ABI3_FORWARD_COMPATIBILITY=1

    或使用预编译 wheel：

    pip install tiktoken --only-binary :all:

  • 离线使用：

  • 首次运行需要互联网下载编码文件（如 cl100k_base.tiktoken）。
  • 可通过缓存文件支持离线模式： 
    import os
os.environ["TIKTOKEN_CACHE_DIR"] = "/path/to/cache"
import tiktoken
tiktoken.get_encoding("cl100k_base")

    • 3. 核心功能与用法

    tiktoken 的核心是通过 Encoding 对象进行文本编码和解码，支持多种编码方式（如 cl100k_base、o200k_base）。以下是主要功能和示例。

    3.1 基本编码和解码

    将文本编码为 token 列表，或将 token 解码为文本。

    import tiktoken

# 获取编码
enc = tiktoken.get_encoding("o200k_base")

# 编码
tokens = enc.encode("hello world")
print(tokens)  # 输出: [24912, 2375]

# 解码
text = enc.decode(tokens)
print(text)  # 输出: hello world

# 验证编码解码一致性
assert enc.decode(enc.encode("hello world")) == "hello world"

    说明：

  • get_encoding：加载指定编码（如 o200k_base）。
  • encode：将文本转换为 token ID 列表。
  • decode：将 token ID 列表转换回文本。
    • 3.2 模型特定编码

    根据模型名称自动选择合适的编码。

    import tiktoken

# 获取 GPT-4o 的编码
enc = tiktoken.encoding_for_model("gpt-4o")
print(enc)  # 输出: <Encoding 'o200k_base'>

# 编码和计数
text = "tiktoken is great!"
tokens = enc.encode(text)
print(len(tokens))  # 输出: 6

    说明：

  • encoding_for_model：根据模型名称（如 gpt-4o、gpt-3.5-turbo）返回正确编码。
  • 支持的模型和编码映射：
  • gpt-4o, o1: o200k_base
  • gpt-3.5-turbo, gpt-4, text-embedding-3-small: cl100k_base
  • 旧模型（如 text-davinci-003）: p50k_base 或 r50k_base
    • 3.3 token 计数

    计算文本的 token 数量，用于验证输入长度或估算成本。

    def num_tokens_from_string(text: str, encoding_name: str) -> int:
    encoding = tiktoken.get_encoding(encoding_name)
    return len(encoding.encode(text))

# 示例
text = "tiktoken is great!"
count = num_tokens_from_string(text, "o200k_base")
print(count)  # 输出: 6

    说明：

  • token 数量用于检查是否超过模型限制（如 gpt-3.5-turbo 的 16k token 限制）。
  • 可估算 API 成本（例如，GPT-3.5-turbo 按 $0.002/千 token 计费）。
    • 3.4 处理对话 token

    计算对话消息的 token 数量，适用于聊天模型。

    def num_tokens_from_messages(messages, model):
    encoding = tiktoken.encoding_for_model(model)
    num_tokens = 0
    for message in messages:
        num_tokens += 4  # 每条消息的固定开销
        for key, value in message.items():
            num_tokens += len(encoding.encode(str(value)))
            if key == "name":
                num_tokens += 1  # 角色名称额外 token
    num_tokens += 2  # 对话结束标记
    return num_tokens

# 示例
messages = [
    {"role": "user", "content": "Hello, how are you?"},
    {"role": "assistant", "content": "I'm doing great, thanks!"}
]
count = num_tokens_from_messages(messages, "gpt-3.5-turbo")
print(count)  # 输出: 约 20（具体数量取决于编码）

    说明：

  • 聊天模型（如 gpt-3.5-turbo）的消息格式包含额外 token（如角色标记）。
  • 确保对话总 token 数不超过模型限制。
    • 3.5 扩展编码

    支持自定义编码，扩展 tiktoken 以支持新模型。

    from tiktoken import Encoding

# 自定义编码
enc = Encoding(
    name="custom_encoding",
    pat_str=r"[^\r\n]+",  # 正则表达式
    mergeable_ranks={"hello": 1, "world": 2},  # BPE 词汇表
    special_tokens={"<|endoftext|>": 1000}
)

tokens = enc.encode("hello world")
print(tokens)  # 输出: [1, 2]

    说明：

  • 通过 tiktoken_ext 插件机制注册自定义编码。
  • 适合需要特定分词规则的场景。
    • 3.6 教育模块

    使用 _educational 子模块学习 BPE 分词。

    from tiktoken._educational import *

# 训练简单 BPE 分词器
enc = train_simple_encoding()
print(enc.encode("hello world"))  # 输出 token ID 列表

# 可视化 GPT-4 编码
enc = SimpleBytePairEncoding.from_tiktoken("cl100k_base")
print(enc.encode("hello world aaaaaaaaaaaa"))

    说明：

  • _educational 模块提供简化的 BPE 实现，帮助理解分词原理。
  • 适合教学或调试。

    • 4. 性能与特点

  • 高效性：3-6 倍于其他开源分词器（如 GPT2TokenizerFast），处理 1GB 文本速度显著。
  • 灵活性：支持多种编码（r50k_base、p50k_base、cl100k_base、o200k_base），适配不同模型。
  • 跨语言：BPE 算法可处理任意字符串，适合多语言文本。
  • 局限性：
  • 首次运行需下载编码文件，可能不适合无网络环境（可通过缓存解决）。
  • 依赖 Rust 编译器，安装可能需额外配置。
  • 主要为 OpenAI 模型设计，适配非 OpenAI 模型需自定义编码。

    • 5. 实际应用场景

  • API 输入验证：检查文本是否超过模型 token 限制，避免 API 错误。
  • 成本管理：估算 OpenAI API 调用费用，优化预算。
  • 文本预处理：为 NLP 任务准备 token 化的输入。
  • 对话管理：处理聊天模型的多轮对话，控制 token 消耗。
  • 教育研究：学习 BPE 分词算法，分析模型输入。

    • 示例（成本估算）：

    import tiktoken

def estimate_api_cost(text, model, price_per_1000_tokens=0.002):
    encoding = tiktoken.encoding_for_model(model)
    token_count = len(encoding.encode(text))
    cost = (token_count / 1000) * price_per_1000_tokens
    return cost

text = "This is a sample text for GPT-3.5-turbo."
cost = estimate_api_cost(text, "gpt-3.5-turbo")
print(f"Estimated cost: ${cost:.4f}")  # 输出: Estimated cost: $0.0001

    说明：

  • 计算文本的 token 数量并估算成本。
  • 帮助开发者控制 API 费用。

    • 示例（Azure OpenAI 集成）：

    import tiktoken
import os
import openai

openai.api_type = "azure"
openai.api_version = "2023-03-15-preview"
openai.api_base = os.getenv("OPENAI_API_BASE")
openai.api_key = os.getenv("OPENAI_API_KEY")

def check_token_limit(text, model, max_tokens=4096):
    encoding = tiktoken.encoding_for_model(model)
    token_count = len(encoding.encode(text))
    return token_count <= max_tokens

# 示例
text = "Hello, this is a test."
if check_token_limit(text, "gpt-3.5-turbo"):
    response = openai.ChatCompletion.create(
        engine="gpt-3.5-turbo",
        messages=[{"role": "user", "content": text}]
    )
    print(response.choices[0].message.content)
else:
    print("Text exceeds token limit")

    说明：

  • 验证输入文本是否符合 Azure OpenAI 模型的 token 限制。
  • 确保 API 调用成功。

    • 6. 与 Streamlit 集成

    以下是一个示例，展示如何在 Streamlit 应用中使用 tiktoken 构建 token 计数工具。

    import streamlit as st
import tiktoken

st.title("Token Counter with Tiktoken")

# 用户输入
text = st.text_area("Enter text to tokenize")
model = st.selectbox("Select model", ["gpt-4o", "gpt-3.5-turbo", "text-embedding-3-small"])

# 计算 token
if text and model:
    encoding = tiktoken.encoding_for_model(model)
    tokens = encoding.encode(text)
    token_count = len(tokens)
    
    st.write(f"**Token Count**: {token_count}")
    st.write(f"**Tokens**: {tokens}")
    st.write(f"**Decoded Text**: {encoding.decode(tokens)}")
    
    # 估算成本
    price_per_1000 = 0.002  # 示例价格（美元）
    cost = (token_count / 1000) * price_per_1000
    st.write(f"**Estimated API Cost**: ${cost:.4f}")

    说明：

  • 用户输入文本和模型，选择编码方式。
  • 显示 token 数量、token ID 列表、解码文本和估算成本。
  • 利用 Streamlit 的交互性，提供直观的用户界面。

    • 7. 注意事项

  • 编码选择：
  • 确保使用与模型匹配的编码（如 gpt-4o 使用 o200k_base）。
  • 错误编码可能导致 token 计数不准确。
  • 离线模式：
  • 缓存编码文件以支持离线使用： 
    RUN python -c "import tiktoken; tiktoken.get_encoding('cl100k_base')"
  • 设置 TIKTOKEN_CACHE_DIR 环境变量。
  • 性能：
  • tiktoken 高效处理大文本，但超大输入可能需分块处理。
  • 避免在循环中重复加载编码对象。
  • 安装问题：
  • 若遇到 Rust 相关错误，安装 Rust 编译器或使用预编译 wheel。
  • 升级 pip 确保安装最新版本： 
    pip install --upgrade pip
  • 解码警告：
  • 单个 token 的 decode 可能因非 UTF-8 边界而丢失信息，建议批量解码。

    • 8. 综合示例

    以下是一个综合示例，展示 token 计数、成本估算和对话 token 处理：

    import tiktoken
import streamlit as st

st.title("Advanced Tiktoken Demo")

# 选择模型
model = st.selectbox("Select Model", ["gpt-4o", "gpt-3.5-turbo"])

# 输入文本
text = st.text_area("Enter Text", "Hello, this is a test for tiktoken.")

# 计算 token
encoding = tiktoken.encoding_for_model(model)
tokens = encoding.encode(text)
token_count = len(tokens)

st.write(f"**Encoding**: {encoding.name}")
st.write(f"**Token Count**: {token_count}")
st.write(f"**Tokens**: {tokens}")

# 估算成本
price_per_1000 = st.number_input("Price per 1000 tokens ($)", value=0.002)
cost = (token_count / 1000) * price_per_1000
st.write(f"**Estimated Cost**: ${cost:.4f}")

# 对话 token 计数
st.subheader("Chat Token Counter")
messages = [
    {"role": "system", "content": st.text_input("System Prompt", "You are a helpful assistant.")},
    {"role": "user", "content": st.text_area("User Message", "Tell me about Python.")}
]

def num_tokens_from_messages(messages, model):
    encoding = tiktoken.encoding_for_model(model)
    num_tokens = 0
    for message in messages:
        num_tokens += 4  # 固定开销
        for key, value in message.items():
            num_tokens += len(encoding.encode(str(value)))
            if key == "name":
                num_tokens += 1
    num_tokens += 2  # 结束标记
    return num_tokens

if st.button("Calculate Chat Tokens"):
    chat_tokens = num_tokens_from_messages(messages, model)
    st.write(f"**Chat Token Count**: {chat_tokens}")
    st.write(f"**Estimated Chat Cost**: ${(chat_tokens / 1000) * price_per_1000:.4f}")

    说明：

  • 支持单文本和对话的 token 计数。
  • 提供成本估算和编码信息。
  • 使用 Streamlit 提供交互式界面。

    • 9. 资源与文档

  • 官方文档：https://github.com/openai/tiktoken
  • PyPI 页面：https://pypi.org/project/tiktoken/
  • OpenAI Cookbook：https://cookbook.openai.com/ （示例代码）
  • 教程：
  • DataCamp Tiktoken 教程：https://www.datacamp.com/tutorial/tiktoken 🔗
  • Estimating GPT Costs：https://www.datacamp.com/tutorial/estimating-cost-gpt-using-tiktoken-python 🔗
  • 社区支持：
  • GitHub Issues：https://github.com/openai/tiktoken/issues
  • OpenAI 开发者社区：https://community.openai.com/ 🔗

    • 作者：彬彬侠

    物联沃分享整理
    物联沃-IOTWORD物联网 » 【Python】tiktoken库：OpenAI开发的BPE分词器库详解与快速字节对编码实践

    代码收藏家 普通

    分享到：

    相关推荐

    发表回复