算法 - 本地部署使用 miniCpmV2-6 、chatglm 这些本地 LLM 的时候，如何实现统计 token？

Question

使用智谱的 glm 在线服务的时候，会返回使用的 token：

• prompt_tokens
• completion_tokens
• total_tokens

Completion(model='glm-4v-plus', created=1731402350, choices=[CompletionChoice(index=0, finish_reason='stop', message=CompletionMessage(content='', role='assistant', tool_calls=None))], request_id='20241112170540c45eaf5a3144487f', id='20241112170540c45eaf5a3144487f', usage=CompletionUsage(prompt_tokens=0, completion_tokens=0, total_tokens=0))

但是如果我是自己本地部署 chatglm、minicpm、qwen 这些模型，怎么统计 token 呢？

示例代码

# test.py
# test.py
import torch
from PIL import Image
from modelscope import AutoModel, AutoTokenizer

model = AutoModel.from_pretrained('OpenBMB/MiniCPM-V-2_6', trust_remote_code=True,
    attn_implementation='sdpa', torch_dtype=torch.bfloat16) # sdpa or flash_attention_2, no eager
model = model.eval().cuda()
tokenizer = AutoTokenizer.from_pretrained('OpenBMB/MiniCPM-V-2_6', trust_remote_code=True)

image = Image.open('image.png').convert('RGB')
question = 'What is in the image?'
msgs = [{'role': 'user', 'content': [image, question]}]

res = model.chat(
    image=None,
    msgs=msgs,
    tokenizer=tokenizer
)
print(res)

## if you want to use streaming, please make sure sampling=True and stream=True
## the model.chat will return a generator
res = model.chat(
    image=None,
    msgs=msgs,
    tokenizer=tokenizer,
    sampling=True,
    stream=True
)

generated_text = ""
for new_text in res:
    generated_text += new_text
    print(new_text, flush=True, end='')

返回的 res 本身就是字符串了，而不是一个结构化对象

这个统计 token 的方式，每个 llm 都一样吗？

Answer 1

回答

在本地部署使用 miniCPM-V2-6、chatglm 等本地 LLM（大型语言模型）时，统计 token 的方式通常并不直接由模型库提供，因为不同的模型库和框架可能有不同的实现方式。以下是一种通用的方法来实现 token 统计：

1. 使用 Tokenizer 统计：
   你可以使用对应的 tokenizer 来统计输入和输出文本的 token 数量。通常，tokenizer 会有一个 encode 或 tokenize 方法，可以将文本转换为 token 列表，然后你可以计算这个列表的长度来得到 token 数量。
2. 修改模型调用代码：
   在你的代码中，你可以在将输入传递给模型之前和从模型接收输出之后，分别使用 tokenizer 来统计 token 数量。

3. 注意事项：

   • 对于图像输入（如你的示例中的 image），tokenizer 通常无法直接处理，你需要先处理图像（如提取文本或使用某种表示方法），然后再统计 token。
   • 对于多模态模型（如同时处理文本和图像的模型），统计 token 可能需要更复杂的逻辑，因为输出可能包含多种类型的数据。
4. 示例修改：
   以下是对你的示例代码的一个简单修改，用于统计输入和输出文本的 token 数量（注意，这只是一个基本示例，可能需要根据你的具体模型和 tokenizer 进行调整）：

import torch
from PIL import Image
from modelscope import AutoModel, AutoTokenizer

model = AutoModel.from_pretrained('OpenBMB/MiniCPM-V-2_6', trust_remote_code=True,
    attn_implementation='sdpa', torch_dtype=torch.bfloat16)
model = model.eval().cuda()
tokenizer = AutoTokenizer.from_pretrained('OpenBMB/MiniCPM-V-2_6', trust_remote_code=True)

image = Image.open('image.png').convert('RGB')
question = 'What is in the image?'
msgs = [{'role': 'user', 'content': [image, question]}]  # 注意：这里 image 可能需要特殊处理才能用于 tokenizer

# 假设我们暂时忽略图像，只统计文本的 token
input_text_tokens = tokenizer.tokenize(question)
prompt_tokens = len(input_text_tokens)

# 由于你的模型调用方式可能不直接返回 tokenized 文本，我们需要从生成的字符串中重新 tokenize
# 注意：这通常不是最佳实践，因为重新 tokenize 可能引入误差，但在这里用于演示
res_text = model.chat(image=None, msgs=msgs, tokenizer=tokenizer)[0]['content']  # 假设返回的是结构化对象，且 'content' 包含生成的文本
output_text_tokens = tokenizer.tokenize(res_text)
completion_tokens = len(output_text_tokens)
total_tokens = prompt_tokens + completion_tokens

print(f"Prompt tokens: {prompt_tokens}")
print(f"Completion tokens: {completion_tokens}")
print(f"Total tokens: {total_tokens}")

注意：上述代码中的 model.chat 调用和返回值的处理可能需要根据你实际使用的模型库进行调整。此外，对于图像输入的处理，你可能需要先将图像转换为文本（如使用 OCR 技术），然后再进行 token 统计。

每个 LLM 的统计 token 方式可能不完全相同，具体取决于其实现和 tokenizer 的设计。因此，在实际应用中，你可能需要参考特定模型库的文档或源代码来了解如何正确地统计 token。