Почему расходится подсчёт токенов в BGE‑M3 и как это исправить в llama_index

Разбираем, почему токены в BGE‑M3 расходятся со счётчиком llama_index и как подключить токенайзер Hugging Face к TokenCountingHandler для точного подсчёта.

Согласование подсчёта токенов между вашим токенайзером и конвейером эмбеддингов напрямую влияет на батчинг, разбиение на фрагменты и оценку стоимости. Если числа не совпадают, вы либо слишком агрессивно обрезаете, либо превышаете лимиты. Частая ловушка — библиотека считает токены другим токенайзером, чем тот, которым реально пользуется модель для эмбеддинга. Ниже — короткое объяснение, почему так происходит с BGE‑M3 в llama_index и как добиться一致ных подсчётов, не запуская сам шаг эмбеддинга.

Воспроизводим расхождение

Следующий фрагмент демонстрирует разницу в количестве токенов между счётчиком токенов llama_index и токенайзером Hugging Face для BGE‑M3. Он использует локальный кэш, если он есть (./embeddings), иначе загружает модель.

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

# Пример текста
sample_text = "Random words. This is a test! A very exciting test, indeed."
# Длина фрагмента
max_chunk = 512

# Создать или загрузить бэкенд для эмбеддингов
def build_encoder(_limit=None):
    print("initializing embeddings...")
    if os.path.exists('./embeddings/models--BAAI--bge-m3'):
        cache_dir = f"./embeddings/models--BAAI--bge-m3/snapshots/{os.listdir('./embeddings/models--BAAI--bge-m3/snapshots')[0]}"
        encoder = HuggingFaceEmbedding(model_name=cache_dir)
    else:
        os.makedirs("./embeddings", exist_ok=True)
        repo_name = "BAAI/bge-m3"
        encoder = HuggingFaceEmbedding(
            model_name=repo_name,
            max_length=_limit,
            cache_folder='./embeddings'
        )
    print("embeddings ready")
    return encoder

# Инициализируем энкодер
embedding_backend = build_encoder(_limit=max_chunk)

# Счётчик токенов с настройками по умолчанию
count_hook = TokenCountingHandler()
hooks = CallbackManager([count_hook])

Settings.embed_model = embedding_backend
Settings.callback_manager = hooks

# Получаем эмбеддинг и считаем токены через callback
_ = Settings.embed_model.get_text_embedding(sample_text)
embed_tok_count = count_hook.total_embedding_token_count

# Считаем токены через токенайзер HF
model_ref = "BAAI/bge-m3"
hf_tok = AutoTokenizer.from_pretrained(model_ref)
encoded = hf_tok(sample_text)
hf_tok_count = len(encoded["input_ids"])

print(f"Original text: {sample_text}")
print(f"Embedding pipeline token count: {embed_tok_count}")
print(f"HF tokenizer token count: {hf_tok_count}")

С такой конфигурацией вы увидите, что числа расходятся.

Что на самом деле происходит

Корневая проблема в том, что подсчёт токенов в llama_index использует токенайзер, отличный от токенайзера вашей модели, если вы явно его не переопределили. По умолчанию применяется токенайзер на базе tiktoken. Посмотреть, что настроено в llama_index во время выполнения, можно так:

from llama_index.core import Settings

print(Settings.tokenizer)

Вывод показывает энкодер tiktoken:

functools.partial(<bound method Encoding.encode of <Encoding 'cl100k_base'>>, allowed_special='all')

Тем временем BGE‑M3 использует токенайзер Hugging Face. AutoTokenizer — это фабрика, и для BGE‑M3 он разворачивается в XLMRobertaTokenizer. Различия в правилах токенизации и объясняют расхождение в счётах.

Решение: используйте токенайзер модели для подсчёта

Чтобы получить совпадающие значения без запуска шага эмбеддинга, подключите токенайзер модели к TokenCountingHandler в llama_index. Есть одна тонкость: TokenCounter в llama_index ожидает, что токенайзер вернёт список input_ids, тогда как токенайзеры Hugging Face возвращают объект BatchEncoding. Небольшая «прокладка» решает это, возвращая только input_ids.

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
from transformers import XLMRobertaTokenizerFast

# Текст и модель
sample_text = "Random words. This is a test! A very exciting test, indeed."
max_chunk = 512
model_ref = "BAAI/bge-m3"

# Адаптер, чтобы счётчик llama_index получал список input_ids
class LlamaIndexTokenizerShim(XLMRobertaTokenizerFast):
    def __call__(self, *args, **kwargs):
        return super().__call__(*args, **kwargs).input_ids

li_tokenizer = LlamaIndexTokenizerShim.from_pretrained(model_ref)

# Инициализируем энкодер HF
embedder = HuggingFaceEmbedding(model_name=model_ref, max_length=max_chunk)

# Подключаем корректный токенайзер к счётчику токенов
counter = TokenCountingHandler(tokenizer=li_tokenizer)
manager = CallbackManager([counter])

Settings.embed_model = embedder
Settings.callback_manager = manager

# Запускаем подсчёт через вызов эмбеддинга (не полагаемся на токенайзер по умолчанию)
_ = Settings.embed_model.get_text_embedding(sample_text)
li_count = counter.total_embedding_token_count

# Для проверки считаем напрямую токенайзером HF
hf_tokenizer = AutoTokenizer.from_pretrained(model_ref)
ref_count = len(hf_tokenizer(sample_text).input_ids)

print(f"Original text: {sample_text}")
print(f"Embedding pipeline token count: {li_count}")
print(f"HF tokenizer token count: {ref_count}")

С этим изменением оба числа совпадают.

Почему это важно

Согласованный подсчёт токенов критичен, когда вы планируете вычислительные ресурсы, режете документы на фрагменты или соблюдаете ограничения max_length. Если счётчик и модель «несогласны», логика батчинга становится ненадёжной: это ведёт к лишним вызовам, неожиданным обрезкам или некорректной отчётности по использованию.

Выводы

Всегда явно задавайте токенайзер при подсчёте токенов в пайплайне эмбеддингов. Если инструменты по умолчанию используют другой токенайзер, чем модель, подключите правильный и при необходимости адаптируйте тип возвращаемого значения. Когда обе стороны «говорят» на одном диалекте токенизации, подсчёты совпадают, а последующая логика остаётся предсказуемой.