如何在嵌入前匹配BGE-M3嵌入模型的Token计数？

我在项目中需要提前知晓BGE-M3嵌入模型拆分文本的Token数量，目前通过嵌入文本后用计数器获取的方式可行，但处理大量文本时计算量太大，想避免这种方法。直接用AutoTokenizer.from_pretrained("BAAI/bge-m3")分词得到的Token数和嵌入模型的计数不一致，推测是模型嵌入前有特殊预处理逻辑，但没找到具体细节。

复现代码如下：

from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_index.core import Settings
from llama_index.core.callbacks import CallbackManager, TokenCountingHandler
from transformers import AutoTokenizer
import os

# 测试文本
text = "Random words. This is a test! A very exciting test, indeed."
chunk_size = 512


# 加载嵌入模型
def create_embedding_model(_chunk_size=None):
    print('loading embeddings...')
    if os.path.exists('./embeddings/models--BAAI--bge-m3'):
        _cache_path = f"./embeddings/models--BAAI--bge-m3/snapshots/{os.listdir('./embeddings/models--BAAI--bge-m3/snapshots')[0]}"
        _embed_model = HuggingFaceEmbedding(model_name=_cache_path)
    else:
        os.makedirs("./embeddings", exist_ok=True)
        _emb_model_name = "BAAI/bge-m3"
        _embed_model = HuggingFaceEmbedding(model_name=_emb_model_name, max_length=_chunk_size,
                                            cache_folder='./embeddings')
    print('embeddings loaded')
    return _embed_model


# 获取嵌入模型
embed_model = create_embedding_model(_chunk_size=chunk_size)

# 创建Token计数处理器
token_counter = TokenCountingHandler()
callback_manager = CallbackManager([token_counter])

Settings.embed_model = embed_model
Settings.callback_manager = callback_manager

# 获取嵌入向量并计数
embedding_vector = Settings.embed_model.get_text_embedding(text)
embedding_tokens = token_counter.total_embedding_token_count

# 直接用Tokenizer分词计数
model_name = "BAAI/bge-m3"
tokenizer = AutoTokenizer.from_pretrained(model_name)

tokenized_text = tokenizer(text)
token_count = len(tokenized_text['input_ids'])

print(f"Original text: {text}")
print(f"The embedding model broke the text into: {embedding_tokens} tokens")
print(f"The tokenizer broke the text into {token_count} tokens")

运行输出：

Original text: Random words. This is a test! A very exciting test, indeed.
The embedding model broke the text into: 15 tokens
The tokenizer broke the text into 18 tokens

请问如何在不执行嵌入操作的前提下，复现BGE-M3嵌入模型使用的Token计数？是否可以按照模型的方式预处理文本，使Tokenizer得到相同的Token数？

问题核心是LlamaIndex的HuggingFaceEmbedding类对文本的预处理逻辑和直接用AutoTokenizer的默认行为不一致，要复现计数需对齐以下步骤：

方法1：直接调用嵌入模型的内部分词方法（推荐） #

无需执行嵌入，直接使用HuggingFaceEmbedding内置的Token计数逻辑，自动对齐预处理规则：

from llama_index.embeddings.huggingface import HuggingFaceEmbedding
import os

text = "Random words. This is a test! A very exciting test, indeed."
chunk_size = 512

# 加载嵌入模型（仅初始化Tokenizer相关逻辑，不加载模型权重）
def get_embedding_tokenizer(_chunk_size=None):
    if os.path.exists('./embeddings/models--BAAI--bge-m3'):
        _cache_path = f"./embeddings/models--BAAI--bge-m3/snapshots/{os.listdir('./embeddings/models--BAAI--bge-m3/snapshots')[0]}"
        embed_model = HuggingFaceEmbedding(model_name=_cache_path, max_length=_chunk_size)
    else:
        os.makedirs("./embeddings", exist_ok=True)
        _emb_model_name = "BAAI/bge-m3"
        embed_model = HuggingFaceEmbedding(model_name=_emb_model_name, max_length=_chunk_size, cache_folder='./embeddings')
    return embed_model

embed_model = get_embedding_tokenizer(_chunk_size=chunk_size)

# 调用内部方法获取Token计数
token_count = embed_model._get_text_token_count(text)

print(f"复现嵌入模型的Token计数：{token_count} tokens")

方法2：手动对齐预处理规则 #

如果不想加载HuggingFaceEmbedding，可手动模拟模型的预处理步骤：

1. 添加BGE-M3检索任务的默认前缀
2. 过滤掉Tokenizer自动添加的<s>（id=0）和</s>（id=2）分隔符

代码示例：

from transformers import AutoTokenizer

model_name = "BAAI/bge-m3"
tokenizer = AutoTokenizer.from_pretrained(model_name)

# BGE-M3检索场景的默认指令前缀
prefix = "Represent this sentence for searching relevant passages: "
processed_text = prefix + text

# 分词并过滤无效Token
tokenized = tokenizer(processed_text)
filtered_input_ids = [id for id in tokenized['input_ids'] if id not in [0, 2]]
token_count = len(filtered_input_ids)

print(f"手动预处理后的Token计数：{token_count} tokens")

以上两种方法得到的Token数都会和嵌入模型的计数一致（均为15）。

内容的提问来源于stack exchange，提问作者ManBearPigeon