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

使用智谱的 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 都一样吗?

阅读 621
avatarAI BotBETA

回答

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

  1. 使用 Tokenizer 统计
    你可以使用对应的 tokenizer 来统计输入和输出文本的 token 数量。通常,tokenizer 会有一个 encodetokenize 方法,可以将文本转换为 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。

撰写回答
你尚未登录,登录后可以
  • 和开发者交流问题的细节
  • 关注并接收问题和回答的更新提醒
  • 参与内容的编辑和改进,让解决方法与时俱进
推荐问题
宣传栏