15 LlamaIndex中的使用模式：预测LLM和嵌入标记计数

LlamaIndex中的使用模式：预测LLM和嵌入标记计数

在LlamaIndex中，为了测量LLM和嵌入的标记计数，你需要执行以下步骤：

设置MockLLM和MockEmbedding对象

首先，你需要设置MockLLM和MockEmbedding对象，这些对象将用于模拟LLM和嵌入调用，以便预测标记使用情况。

from llama_index.core.llms import MockLLM
from llama_index.core import MockEmbedding

llm = MockLLM(max_tokens=256)
embed_model = MockEmbedding(embed_dim=1536)

设置TokenCountingCallback处理程序

接下来，你需要设置TokenCountingCallback处理程序，该处理程序将使用tiktoken库来计算标记数。

import tiktoken
from llama_index.core.callbacks import CallbackManager, TokenCountingHandler

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

callback_manager = CallbackManager([token_counter])

将它们添加到全局设置

然后，将MockLLM、MockEmbedding和CallbackManager添加到全局设置中。

from llama_index.core import Settings

Settings.llm = llm
Settings.embed_model = embed_model
Settings.callback_manager = callback_manager

构建索引

现在，你可以构建索引并加载文档。

from llama_index.core import VectorStoreIndex, SimpleDirectoryReader

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

index = VectorStoreIndex.from_documents(documents)

测量计数

构建索引后，你可以测量嵌入和LLM的标记计数。

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

运行查询，再次测量

最后，你可以运行查询并再次测量标记计数。

query_engine = index.as_query_engine()

response = query_engine.query("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",
)

通过这些步骤，你可以在LlamaIndex中预测LLM和嵌入调用的标记使用情况，从而更好地管理和优化你的LLM应用的成本。