使用OpenAI提供的tiktoken库来计算输入字符串中的token数量

21 min read
pip install tiktoken

然后,您可以使用以下代码计算输入文本的token数量:

pythonCopy code
from tiktoken import Tokenizer
from tiktoken.tokenizer import Tokenizer, TokenizerException

def count_tokens(text):
    tokenizer = Tokenizer()
    token_count = 0
    try:
        tokens = tokenizer.tokenize(text)
        token_count = len(tokens)
    except TokenizerException as e:
        print(f"Error tokenizing text: {e}")
    return token_count

input_text = "这是一个示例文本。"  # 使用您自己的文本替换此处
tokens = count_tokens(input_text)
print(f"Token count: {tokens}")

这个简单的函数将文本作为输入,使用tiktoken库的Tokenizer对其进行分词,并返回token数量。

请注意,对于非英语文本,计算token数量可能会有所不同。在这种情况下,您可能需要对tiktoken库进行调整以更准确地计算token数量。

stream 模式如何统计

import openai
from tiktoken import Tokenizer

openai.api_key = "your-api-key"  # 使用您的API密钥替换

def count_tokens(text):
    tokenizer = Tokenizer()
    tokens = tokenizer.tokenize(text)
    return len(tokens)

def generate_streaming(prompt, tokens_limit):
    model_engine = "text-davinci-002"  # 使用您想要使用的模型
    stream = openai.Completion.create_stream(
        engine=model_engine,
        prompt=prompt,
        max_tokens=tokens_limit,
        n=1,
        stop=None,
        temperature=0.8,
    )

    token_count = 0
    input_token_count = count_tokens(prompt)

    print(f"Input tokens: {input_token_count}")
    print("Generated text:")

    for chunk in stream.iter_chunks():
        token_count += count_tokens(chunk.choices[0].text)
        print(chunk.choices[0].text, end="")
        if token_count >= tokens_limit:
            break

    print(f"\n\nTotal tokens (including input): {input_token_count + token_count}")

prompt_text = "Once upon a time"  # 使用您自己的文本替换
tokens_limit = 100  # 设置生成文本的token上限
generate_streaming(prompt_text, tokens_limit)