使用LlamaIndex进行Token计数的实战指南

在人工智能领域，特别是在自然语言处理（NLP）任务中，理解和跟踪Token的使用情况是非常重要的。这篇文章将介绍如何使用LlamaIndex库来进行Token计数，并提供一些实用的代码示例，以便你在自己的项目中应用这些技术。

环境设置

首先，我们需要设置回调和服务上下文。通过全局设置，我们可以在不需要每次查询时都传递这些设置的情况下使用它们。

import os

os.environ["OPENAI_API_KEY"] = "sk-..."  # 请替换为您的实际API密钥

import tiktoken
from llama_index.core.callbacks import CallbackManager, TokenCountingHandler
from llama_index.llms.openai import OpenAI
from llama_index.core import Settings

token_counter = TokenCountingHandler(
    tokenizer=tiktoken.encoding_for_model("gpt-3.5-turbo").encode
)

Settings.llm = OpenAI(model="gpt-3.5-turbo", temperature=0.2)
Settings.callback_manager = CallbackManager([token_counter])

Token计数

Token计数器将跟踪嵌入、提示和完成的Token使用情况。Token计数是累积的，只有在你选择重置它们时才会重置。

token_counter.reset_counts()

嵌入Token使用情况

现在，设置已经完成，让我们跟踪嵌入Token的使用情况。

!mkdir -p 'data/paul_graham/'
!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/docs/examples/data/paul_graham/paul_graham_essay.txt' -O 'data/paul_graham/paul_graham_essay.txt'

from llama_index.core import SimpleDirectoryReader

documents = SimpleDirectoryReader("./data/paul_graham").load_data()

from llama_index.core import VectorStoreIndex

index = VectorStoreIndex.from_documents(documents)

print(token_counter.total_embedding_token_count)  # 输出总嵌入Token计数

LLM + 嵌入Token使用情况

接下来，让我们测试一个查询并查看计数结果。

query_engine = index.as_query_engine(similarity_top_k=4)
response = query_engine.query("作者在成长过程中做了什么？")

print(
    "Embedding Tokens: ",
    token_counter.total_embedding_token_count,
    "\n",
    "LLM Prompt Tokens: ",
    token_counter.prompt_llm_token_count,
    "\n",
    "LLM Completion Tokens: ",
    token_counter.completion_llm_token_count,
    "\n",
    "Total LLM Token Count: ",
    token_counter.total_llm_token_count,
    "\n",
)

Token计数 + 流式处理

Token计数处理程序还可以在流式处理中处理Token计数。这里，Token计数将在流完成后进行。

token_counter.reset_counts()

query_engine = index.as_query_engine(similarity_top_k=4, streaming=True)
response = query_engine.query("Interleaf发生了什么？")

# 完成流处理
for token in response.response_gen:
    continue

print(
    "Embedding Tokens: ",
    token_counter.total_embedding_token_count,
    "\n",
    "LLM Prompt Tokens: ",
    token_counter.prompt_llm_token_count,
    "\n",
    "LLM Completion Tokens: ",
    token_counter.completion_llm_token_count,
    "\n",
    "Total LLM Token Count: ",
    token_counter.total_llm_token_count,
    "\n",
)

高级用法

Token计数器在一个名为TokenCountingEvent的对象中跟踪每个Token使用事件。这个对象包含以下属性：

• prompt -> 发送到LLM或嵌入模型的提示字符串
• prompt_token_count -> LLM提示的Token计数
• completion -> 从LLM收到的字符串完成（不用于嵌入）
• completion_token_count -> LLM完成的Token计数（不用于嵌入）
• total_token_count -> 事件的总提示+完成Token计数
• event_id -> 事件的字符串ID，与其他回调处理程序对齐

这些事件在Token计数器中以两个列表进行跟踪：

• llm_token_counts
• embedding_token_counts

让我们快速查看这些事件的示例。

print("Num LLM token count events: ", len(token_counter.llm_token_counts))
print(
    "Num Embedding token count events: ",
    len(token_counter.embedding_token_counts),
)

print("prompt: ", token_counter.llm_token_counts[0].prompt[:100], "...\n")
print(
    "prompt token count: ",
    token_counter.llm_token_counts[0].prompt_token_count,
    "\n",
)

print(
    "completion: ", token_counter.llm_token_counts[0].completion[:100], "...\n"
)
print(
    "completion token count: ",
    token_counter.llm_token_counts[0].completion_token_count,
    "\n",
)

print("total token count", token_counter.llm_token_counts[0].total_token_count)

可能遇到的错误

1. API密钥错误: 如果API密钥无效或未设置，代码将无法正常工作。请确保使用有效的API密钥，并正确设置环境变量。
2. 网络问题: 在下载数据或访问API时，可能会遇到网络问题。请确保网络连接正常，并且能够访问所需的资源。
3. Token计数错误: 如果Token计数结果与预期不符，可能是由于设置或编码问题。请确保按照示例代码正确设置了Token计数器。

如果你觉得这篇文章对你有帮助,请点赞,关注我的博客,谢谢!

参考资料:

• LlamaIndex文档