[1338]python tiktoken基本用法

@[toc]

tiktoken 是 OpenAI 开发的一个快速 BPE 分词器库，用于将文本转换成 GPT 模型使用的 token ID。以下是详细的用法指南：

安装

pip install tiktoken

基本用法

1. 初始化编码器

import tiktoken

# 获取不同模型的编码器
enc = tiktoken.get_encoding("cl100k_base")  # GPT-4, GPT-3.5-turbo
enc = tiktoken.encoding_for_model("gpt-4")  # 根据模型名自动选择编码

# 支持的编码类型：
# - "gpt2"：GPT-2
# - "p50k_base"：Codex 模型
# - "cl100k_base"：GPT-4, GPT-3.5-turbo
# - "o200k_base"：最新模型

2. 基本编码解码

import tiktoken

# 初始化编码器
enc = tiktoken.get_encoding("cl100k_base")

# 编码文本 -> token IDs
text = "Hello, world! 你好，世界！"
token_ids = enc.encode(text)
print("Token IDs:", token_ids)
# 输出: [9906, 11, 1917, 0, 123, 123, 123, 123] (示例)

# 解码 token IDs -> 文本
decoded_text = enc.decode(token_ids)
print("Decoded text:", decoded_text)

# 直接查看 tokens
tokens = enc.encode_ordinary(text)
print("Tokens:", tokens)

3. 按模型使用

import tiktoken

def num_tokens_from_string(text: str, model_name: str) -> int:
    """返回文本的 token 数量"""
    encoding = tiktoken.encoding_for_model(model_name)
    return len(encoding.encode(text))

# 示例
text = "这是一个测试句子。"
models = ["gpt-4", "gpt-3.5-turbo", "text-davinci-003"]

for model in models:
    token_count = num_tokens_from_string(text, model)
    print(f"{model}: {token_count} tokens")

高级功能

1. 批量编码和解码

import tiktoken

enc = tiktoken.get_encoding("cl100k_base")

# 批量编码
texts = ["Hello world", "Python编程", "机器学习"]
batch_tokens = [enc.encode(text) for text in texts]
print("Batch tokens:", batch_tokens)

# 批量解码
decoded_texts = [enc.decode(tokens) for tokens in batch_tokens]
print("Decoded texts:", decoded_texts)

2. 处理特殊 tokens

import tiktoken

enc = tiktoken.get_encoding("cl100k_base")

# 允许特殊 tokens
text = "Hello<|endoftext|>World"
tokens_with_special = enc.encode(text, allowed_special="all")
print("With special tokens:", tokens_with_special)

# 不允许特殊 tokens（默认）
tokens_without_special = enc.encode(text, allowed_special={})
print("Without special tokens:", tokens_without_special)

# 自定义允许的特殊 tokens
tokens_custom = enc.encode(text, allowed_special={"<|endoftext|>"})
print("Custom special tokens:", tokens_custom)

3. 分块处理长文本

import tiktoken

def chunk_text_by_tokens(text, model_name, max_tokens=2048):
    """将长文本按 token 数量分块"""
    encoding = tiktoken.encoding_for_model(model_name)
    tokens = encoding.encode(text)
    
    chunks = []
    for i in range(0, len(tokens), max_tokens):
        chunk_tokens = tokens[i:i + max_tokens]
        chunk_text = encoding.decode(chunk_tokens)
        chunks.append(chunk_text)
    
    return chunks

# 使用示例
long_text = "很长很长的文本..." * 1000
chunks = chunk_text_by_tokens(long_text, "gpt-4", max_tokens=1000)
print(f"分成了 {len(chunks)} 个块")

实际应用示例

1. 计算聊天消息的 tokens

import tiktoken

def num_tokens_from_messages(messages, model="gpt-3.5-turbo-0613"):
    """计算聊天消息的总 token 数"""
    try:
        encoding = tiktoken.encoding_for_model(model)
    except KeyError:
        encoding = tiktoken.get_encoding("cl100k_base")
    
    # 不同模型的 token 计算方式
    if model in {
        "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",
    }:
        tokens_per_message = 3
        tokens_per_name = 1
    else:
        tokens_per_message = 3
        tokens_per_name = 1
    
    num_tokens = 0
    for message in messages:
        num_tokens += tokens_per_message
        for key, value in message.items():
            num_tokens += len(encoding.encode(value))
            if key == "name":
                num_tokens += tokens_per_name
    
    num_tokens += 3  # 每个回复都以 assistant 开头
    return num_tokens

# 使用示例
messages = [
    {"role": "system", "content": "你是一个有帮助的助手。"},
    {"role": "user", "content": "你好，请介绍一下Python。"}
]

token_count = num_tokens_from_messages(messages, "gpt-4")
print(f"消息总 tokens: {token_count}")

2. 价格计算器

import tiktoken

class TokenCalculator:
    def __init__(self, model_name):
        self.model_name = model_name
        self.encoding = tiktoken.encoding_for_model(model_name)
        self.prices = {
            "gpt-4": {"input": 0.03, "output": 0.06},  # 每千token价格（美元）
            "gpt-3.5-turbo": {"input": 0.0015, "output": 0.002},
        }
    
    def calculate_cost(self, input_text, output_text=""):
        input_tokens = len(self.encoding.encode(input_text))
        output_tokens = len(self.encoding.encode(output_text))
        
        if self.model_name in self.prices:
            price_info = self.prices[self.model_name]
            cost = (input_tokens / 1000 * price_info["input"] + 
                   output_tokens / 1000 * price_info["output"])
            return {
                "input_tokens": input_tokens,
                "output_tokens": output_tokens,
                "total_tokens": input_tokens + output_tokens,
                "cost_usd": round(cost, 4)
            }
        return None

# 使用示例
calculator = TokenCalculator("gpt-4")
result = calculator.calculate_cost("请写一篇关于AI的文章", "AI是...")
print(result)

3. 文本截断工具

import tiktoken

def truncate_text_by_tokens(text, model_name, max_tokens):
    """根据 token 数量截断文本"""
    encoding = tiktoken.encoding_for_model(model_name)
    tokens = encoding.encode(text)
    
    if len(tokens) <= max_tokens:
        return text
    
    truncated_tokens = tokens[:max_tokens]
    return encoding.decode(truncated_tokens)

# 使用示例
long_text = "这是一段很长的文本..." * 100
truncated = truncate_text_by_tokens(long_text, "gpt-4", max_tokens=50)
print(f"截断后长度: {len(truncated)}")

性能优化技巧

import tiktoken
import time

# 1. 重复使用编码器实例
enc = tiktoken.get_encoding("cl100k_base")  # 初始化一次，多次使用

# 2. 批量处理
texts = ["text1", "text2", "text3"] * 1000

# 慢的方式
start = time.time()
for text in texts:
    tokens = tiktoken.get_encoding("cl100k_base").encode(text)
print(f"单次编码: {time.time() - start:.2f}秒")

# 快的方式
start = time.time()
for text in texts:
    tokens = enc.encode(text)
print(f"复用编码器: {time.time() - start:.2f}秒")

常见问题处理

1. 处理编码错误

import tiktoken

def safe_encode(text, encoding_name="cl100k_base"):
    """安全编码，处理异常字符"""
    enc = tiktoken.get_encoding(encoding_name)
    
    try:
        return enc.encode(text)
    except Exception as e:
        print(f"编码错误: {e}")
        # 替换或忽略无法编码的字符
        text_clean = text.encode('utf-8', 'ignore').decode('utf-8')
        return enc.encode(text_clean)

# 使用示例
problematic_text = "正常文本 + 异常字符: \uFFFF"
tokens = safe_encode(problematic_text)

2. 多语言支持

import tiktoken

def compare_token_count(texts, model_name="gpt-4"):
    """比较不同语言文本的 token 数量"""
    enc = tiktoken.encoding_for_model(model_name)
    
    for text in texts:
        tokens = enc.encode(text)
        print(f"文本: {text[:30]}...")
        print(f"字符数: {len(text)}, Token数: {len(tokens)}")
        print(f"比例: {len(tokens)/len(text):.2f}\n")

# 测试不同语言
texts = [
    "This is an English sentence.",
    "这是一个中文句子。",
    "これは日本語の文章です。",
    "이것은 한국어 문장입니다."
]

compare_token_count(texts)

这些用法涵盖了 tiktoken 的主要功能，可以帮助您有效地处理文本 tokenization 任务。