令牌计数处理器

本笔记本将逐步介绍如何使用 TokenCountingHandler 以及如何利用它来追踪您随时间变化的提示词、完成内容和嵌入标记使用情况。

如果您在 Colab 上打开这个笔记本，您可能需要安装 LlamaIndex 🦙。

%pip install llama-index-llms-openai

!pip install llama-index

设置

在这里，我们设置回调和服务上下文。我们配置全局设置，这样就不必担心将其传递到索引和查询中。

import os

os.environ["OPENAI_API_KEY"] = "sk-..."

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_counter.reset_counts()。

嵌入令牌使用情况

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

下载数据

!mkdir -p 'data/paul_graham/'
!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/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_counter.reset_counts()

LLM + 嵌入令牌使用量

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

query_engine = index.as_query_engine(similarity_top_k=4)
response = query_engine.query("What did the author do growing up?")

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",
)

Embedding Tokens:  8
 LLM Prompt Tokens:  4518
 LLM Completion Tokens:  45
 Total LLM Token Count:  4563

令牌计数 + 流式处理！

令牌计数处理器同样处理流式传输过程中的令牌计数。

在这里，令牌计数只会在流完成时发生。

token_counter.reset_counts()

query_engine = index.as_query_engine(similarity_top_k=4, streaming=True)
response = query_engine.query("What happened at Interleaf?")

# finish the stream
for token in response.response_gen:
    # print(token, end="", flush=True)
    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",
)

Embedding Tokens:  6
 LLM Prompt Tokens:  4563
 LLM Completion Tokens:  123
 Total LLM Token Count:  4686

高级用法

令牌计数器在一个名为 TokenCountingEvent 的对象中追踪每个令牌使用事件。该对象具有以下属性：

prompt -> 发送给LLM或嵌入模型的提示字符串
prompt_token_count -> 大语言模型提示词的分词数量
completion -> 从LLM接收的字符串补全结果（不用于嵌入）
completion_token_count -> 大语言模型补全的令牌计数（不用于嵌入）
total_token_count -> 该事件中提示词与补全词符的总数
event_id -> 事件的字符串ID，与其他回调处理程序保持一致

这些事件在令牌计数器中通过两个列表进行追踪：

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),
)

Num LLM token count events:  2
Num Embedding token count events:  1

这很有道理！之前的查询嵌入了查询文本，然后进行了2次LLM调用（因为top k设置为4，默认分块大小为1024，需要进行两次独立调用以便LLM能够读取所有检索到的文本）。

接下来，让我们快速看看单个事件的这些事件是什么样的。

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)

prompt:  system: You are an expert Q&A system that is trusted around the world.
Always answer the query using ...

prompt token count:  3873

completion:  assistant: At Interleaf, the company had added a scripting language inspired by Emacs and made it a  ...

completion token count:  95

total token count 3968