[Python] How can I match the token count used by BGE-M3 embedding model before embedding?

For my particular project, it would be very helpful to know how many tokens the BGE-M3 embedding model would break a string down into before I embed the text. I could embed the string and count the tokens with using the following code

Settings.callback_manager = callback_manager
embedding_vector = Settings.embed_model.get_text_embedding(text)
embedding_tokens = token_counter.total_embedding_token_count


but unfortunately, embedding large amounts of text is a relatively computationally heavy problem, so I would prefer not to use this method. After doing some digging, I realized that I could use

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

to tokenize text, but oddly enough any text I embed using this method gives different results than using the previously mentioned method. I think the crux of my issue might be that the BGE-M3 embedding model pre-processes text prior to embedding. I have tried to google exactly what this pre-processing step looks like, but have been unable to find it so far. Below is some code that will allow us to re-create the issue I am talking about. Please note that this script assumes you have the model saved in ./embeddings. If you don't, it will attempt to automatically download the model, which is > 2GB.

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

# String to test on
text = "Random words. This is a test! A very exciting test, indeed."
# just set chunk_size to 512
chunk_size = 512


# Load in the embedding model. If it does not exist, go ahead and download it
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


# Grab embedding model
embed_model = create_embedding_model(_chunk_size=chunk_size)

# Create a token counting handler
token_counter = TokenCountingHandler()
callback_manager = CallbackManager([token_counter])

Settings.embed_model = embed_model
Settings.callback_manager = callback_manager

# Grab the embedding vector from the model
embedding_vector = Settings.embed_model.get_text_embedding(text)

# Grab the count of tokens from the embedding vector
embedding_tokens = token_counter.total_embedding_token_count

# Just the 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")



When I run the above, I get the following output

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
​

How can I replicate the token count that the BGE-M3 embedding model uses, without running the embedding itself? Is there a way to pre-process the text in the same way the embedding model does, so that the tokenizer gives the same token count?

Continue reading...