Using custom vocabulary with subword tokenizer

[aoezdTchibo]

The rationale for this problem stems from a related issue in the sentence-transformer project: UKPLab/sentence-transformers#3313

For testing purposes I created following training script (source @tomaarsen comment on UKPLab/sentence-transformers#3281):

Training script with a distilled model with custom vocabulary

import logging
import random

from nltk.corpus import words
from sentence_transformers import (
    SentenceTransformer,
    SentenceTransformerTrainer,
    SentenceTransformerTrainingArguments,
)
from sentence_transformers.losses import MultipleNegativesRankingLoss
from sentence_transformers.models.StaticEmbedding import StaticEmbedding
from sentence_transformers.training_args import BatchSamplers, MultiDatasetBatchSamplers

from datasets import Dataset, DatasetDict

logging.basicConfig(
    format="%(asctime)s - %(message)s", datefmt="%Y-%m-%d %H:%M:%S", level=logging.INFO
)
random.seed(12)


def load_train_dataset():
    try:
        return DatasetDict.load_from_disk("datasets/train_dataset")
    except FileNotFoundError:
        wordlist = words.words()

        dataset_size = 1_000_000
        dataset_names = ["dataset_A", "dataset_B", "dataset_C", "dataset_D", "dataset_E"]
        datasets = {}
        for dataset_name in dataset_names:
            # Create a dummy dataset with random sentences
            sentence_As = [" ".join(random.sample(wordlist, k=30)) for _ in range(dataset_size)]
            sentence_Bs = [" ".join(random.sample(wordlist, k=30)) for _ in range(dataset_size)]
            dataset = Dataset.from_dict({"sentence_A": sentence_As, "sentence_B": sentence_Bs})
            datasets[dataset_name] = dataset

        train_dataset = DatasetDict(datasets)
        train_dataset.save_to_disk("datasets/train_dataset")

        return train_dataset


def main():
    base_model_name = "intfloat/multilingual-e5-large"
    base_model = SentenceTransformer(base_model_name)
    new_tokens = ["my_custom_token"]
    custom_vocab = list(base_model.tokenizer.get_vocab()) + new_tokens
    static_embedding = StaticEmbedding.from_distillation(
        base_model_name,
        device="cpu",
        vocabulary=custom_vocab,
    )
    model = SentenceTransformer(
        modules=[static_embedding],
        device="cpu",
    )
    train_dataset = load_train_dataset()
    loss = MultipleNegativesRankingLoss(
        model,
        scale=50,
    )
    run_name = f"{base_model_name}_custom_vocab"
    args = SentenceTransformerTrainingArguments(
        # Required parameter:
        output_dir=f"models/{run_name}",
        # Optional training parameters:
        num_train_epochs=1,
        per_device_train_batch_size=2048,
        per_device_eval_batch_size=2048,
        learning_rate=2e-1,
        warmup_ratio=0.1,
        fp16=False,  # Set to False if you get an error that your GPU can't run on FP16
        bf16=True,  # Set to True if you have a GPU that supports BF16
        batch_sampler=BatchSamplers.NO_DUPLICATES,
        # MultipleNegativesRankingLoss benefits from no duplicate samples in a batch
        multi_dataset_batch_sampler=MultiDatasetBatchSamplers.PROPORTIONAL,
        # Optional tracking/debugging parameters:
        eval_strategy="no",
        save_strategy="no",
        logging_steps=50,
        logging_first_step=True,
    )

    # 6. Create a trainer & train
    trainer = SentenceTransformerTrainer(
        model=model,
        args=args,
        train_dataset=train_dataset,
        loss=loss,
    )
    print("train...")
    trainer.train()

    # 7. Save the trained model
    model.save_pretrained(f"models/{run_name}")

if __name__ == "__main__":
    main()

If I run my script I get the following error (ValueError: Number of tokens (250002) does not match number of vectors (404627)):

2025-04-14 10:19:18 - Use pytorch device_name: mps
2025-04-14 10:19:18 - Load pretrained SentenceTransformer: intfloat/multilingual-e5-large
2025-04-14 10:19:25 - The `apply_zipf` parameter is deprecated and will be removed in the next release. Zipf weighting is applied based on the sif_coefficient parameter. If this is set to None, no weighting is applied.
2025-04-14 10:19:26 - Removed 95373 duplicate tokens.
2025-04-14 10:19:26 - Removed 4 multiword tokens.
2025-04-14 10:19:26 - Adding 154628 tokens to the vocabulary. Removed 95377 tokens during preprocessing.
100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 396/396 [02:50<00:00,  2.32it/s]
2025-04-14 10:22:26 - Applying PCA with n_components 256
2025-04-14 10:22:40 - Reduced dimensionality from 1024 to 256.
2025-04-14 10:22:40 - Explained variance ratio: 0.914.
2025-04-14 10:22:40 - Explained variance: 167.012.
2025-04-14 10:22:40 - Estimating word frequencies using Zipf's law, and then applying SIF.
Traceback (most recent call last):
  File "(...)/train.py", line 99, in <module>
    main()
  File "(...)/train.py", line 49, in main
    static_embedding = StaticEmbedding.from_distillation(
                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "(...)/.venv/lib/python3.11/site-packages/sentence_transformers/models/StaticEmbedding.py", line 175, in from_distillation
    static_model = distill(
                   ^^^^^^^^
  File "(...)/.venv/lib/python3.11/site-packages/model2vec/distill/distillation.py", line 241, in distill
    return distill_from_model(
           ^^^^^^^^^^^^^^^^^^^
  File "(...)/.venv/lib/python3.11/site-packages/model2vec/distill/distillation.py", line 136, in distill_from_model
    return StaticModel(
           ^^^^^^^^^^^^
  File "(...)/.venv/lib/python3.11/site-packages/model2vec/model.py", line 50, in __init__
    raise ValueError(f"Number of tokens ({len(tokens)}) does not match number of vectors ({vectors.shape[0]})")
ValueError: Number of tokens (250002) does not match number of vectors (404627)

What issues exist with the script, or how can I effectively expand the vocabulary to include my required custom tokens?

Given that the subword tokenizer feature hasn't been released yet, when can we expect its availability? This information is crucial as it impacts our own release schedule. 😁

[Pringled]

Hi @aoezdTchibo, thanks for opening this issue. We plan on releasing this feature soon (within a week) after a bit more testing. For now, would it work for you to install the main branch directly?

[aoezdTchibo]

Hi @Pringled!

Thanks for the update!

Yes, it would definitely be okay to work with the main branch. The internal merge request can wait until you make the release. That's not a problem. Nevertheless, I have reinstalled the main branch locally and executed the script mentioned above. I see new content in the logs that I hadn't seen before (-> working on main). Unfortunately, I'm still encountering the error: ValueError: Number of tokens (250002) does not match number of vectors (404627)

[Pringled]

Hi @aoezdTchibo , I see the issue now. You are doing the following:

    new_tokens = ["my_custom_token"]
    custom_vocab = list(base_model.tokenizer.get_vocab()) + new_tokens

But you only need to pass the new tokens to vocabulary, the base vocabulary from the sentence transformer is always used. The following works for me for example:

from model2vec.distill import distill

base_model_name = "intfloat/multilingual-e5-large"

m2v_model = distill(model_name=base_model_name, vocabulary=["my_custom_token"])

Let me know if this fixes your issue.

[tomaarsen]

My apologies, I proposed list(model[0].tokenizer.get_vocab()) + ["my_special_token", "my_special_token2"] in my response on the Sentence Transformers side, I didn't realise you only had to add new tokens.

Perhaps it makes sense to remove all tokens from the vocabulary that are already in the tokenizer vocabulary in model2vec, though? That way you won't get crashes if you (accidentally) add a token to the vocabulary that's already in there.

• Tom Aarsen

[Pringled]

We actually do that already (removing duplicates), so this might be a bug with the new supported backends. I'll investigate this, thanks for the extra context @tomaarsen!

[aoezdTchibo]

With the help of @Pringled's comments, I was able to successfully execute the script. The command line output also shows that duplicates are filtered out if they exist.

But when we integrated our list of brand and product names as an extended vocabulary, the following issue arose:

from sentence_transformers.models import StaticEmbedding

static_embedding = StaticEmbedding.from_distillation(
    "intfloat/multilingual-e5-small",  # Smaller model for faster experimentation
    device="cpu",
    vocabulary=[
        "esperto", "Esperto",  # ...
    ],
)

# ValueError: Number of tokens (250000) does not match number of vectors (250001)

"Esperto" is one of our product lines and we would like to consider both formats within the vocabulary ("esperto" vs. "Esperto").

EDIT: I had to update example, since the previous example had no problems...

EDIT 2: After some experimenting I can say, that the string "esperto" in general is not processable... 😁

[Pringled]

@aoezdTchibo this is definitely a bug with the newly supported backends, we'll work on a fix. It's basically failing when trying to add tokens that are already in vocab (and "esperto" is in vocab already). So "Esperto" will work, but "esperto" won't in this case. I'll let you know once we have found a fix for this, for now, you can optionally try to only add new tokens that are not already in vocab.

[aoezdTchibo]

@Pringled Thanks for the update!

"esperto" is in vocab already

Is that true? I am not able to confirm that for intfloat/multilingual-e5-small:

from sentence_transformers import SentenceTransformer
from sentence_transformers.models import StaticEmbedding

model_name = "intfloat/multilingual-e5-small"
embed_model = SentenceTransformer(model_name)
print("esperto" in embed_model.tokenizer.get_vocab()). # --> False

Also within the console while distillation I also get the info, that no duplicates are present if "esperto" is in the new tokens: 2025-04-16 14:03:16 - Adding 6 tokens to the vocabulary. Removed 0 tokens during preprocessing.

What am I missing?

[Pringled]

@aoezdTchibo it's actually a subword ("_esperto") which I think is related to the bug. The vocab is a dict btw, the keys are the actual tokens. So the following will evaluate to True:

from sentence_transformers import SentenceTransformer

model_name = "intfloat/multilingual-e5-small"
embed_model = SentenceTransformer(model_name)
vocab = list(embed_model.tokenizer.get_vocab().keys())
print("▁esperto" in vocab) # --> True

[stephantul]

@aoezdTchibo This is now fixed on main. Let us know if it works, we're releasing "soon".

[aoezdTchibo]

@stephantul My local test with the token "esperto" is now running smoothly, and therefore, I can expand the vocabulary as needed. I am closing this issue.

Fantastic work that helps us immensely. Thank you for your efforts!

[aoezdTchibo]

Unfortunately, the problem is not completely solved. After using our automatically generated list of brand names as additional vocabulary, I found more examples where the distillation didn't work. 😢

Here are some examples of tokens that didn't work: ["ELLE", "RIN", "ADE", "weka", "ONE"]
These are examples which are already in the tokenizers vocabulary. But other already existing tokens like ["LG", "Samsung", "Galaxy", "Note", "Apple"] do not trigger the error. So I am unable to see a specific pattern.

StaticEmbedding.from_distillation(
    base_model_name,
    device="cpu",
    vocabulary=[
        "ELLE",
    ],
)

The `apply_zipf` parameter is deprecated and will be removed in the next release. Zipf weighting is applied based on the sif_coefficient parameter. If this is set to None, no weighting is applied.
100%|██████████| 245/245 [00:11<00:00, 20.44it/s]
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[39], line 1
----> 1 StaticEmbedding.from_distillation(
      2     base_model_name,
      3     device="cpu",
      4     vocabulary=[
      5         "ELLE",
      6     ],
      7 )

File (...)/sentence_transformers/models/StaticEmbedding.py:175, in StaticEmbedding.from_distillation(cls, model_name, vocabulary, device, pca_dims, apply_zipf, use_subword)
    170     raise ImportError(
    171         "To use this method, please install the `model2vec` package: `pip install model2vec[distill]`"
    172     )
    174 device = get_device_name()
--> 175 static_model = distill(
    176     model_name,
    177     vocabulary=vocabulary,
    178     device=device,
    179     pca_dims=pca_dims,
    180     apply_zipf=apply_zipf,
    181     use_subword=use_subword,
    182 )
    183 if isinstance(static_model.embedding, np.ndarray):
    184     embedding_weights = torch.from_numpy(static_model.embedding)

File (...)/model2vec/distill/distillation.py:250, in distill(model_name, vocabulary, device, pca_dims, apply_zipf, sif_coefficient, use_subword, token_remove_pattern, trust_remote_code, quantize_to)
    247 model: PreTrainedModel = AutoModel.from_pretrained(model_name, trust_remote_code=trust_remote_code)
    248 tokenizer: PreTrainedTokenizerFast = AutoTokenizer.from_pretrained(model_name, trust_remote_code=trust_remote_code)
--> 250 return distill_from_model(
    251     model=model,
    252     tokenizer=tokenizer,
    253     vocabulary=vocabulary,
    254     device=device,
    255     pca_dims=pca_dims,
    256     apply_zipf=apply_zipf,
    257     use_subword=use_subword,
    258     token_remove_pattern=token_remove_pattern,
    259     sif_coefficient=sif_coefficient,
    260     quantize_to=quantize_to,
    261 )

File (...)/model2vec/distill/distillation.py:143, in distill_from_model(model, tokenizer, vocabulary, device, pca_dims, apply_zipf, sif_coefficient, use_subword, token_remove_pattern, quantize_to)
    140         logger.info("No model info found for the model. Setting language to None.")
    141         language = None
--> 143 return StaticModel(
    144     vectors=embeddings,
    145     tokenizer=backend_tokenizer,
    146     config=config,
    147     base_model_name=model_name,
    148     language=language,
    149     normalize=True,
    150 )

File (...)/model2vec/model.py:51, in StaticModel.__init__(self, vectors, tokenizer, config, normalize, base_model_name, language)
     48 self.embedding = vectors
     50 if len(tokens) != vectors.shape[0]:
---> 51     raise ValueError(f"Number of tokens ({len(tokens)}) does not match number of vectors ({vectors.shape[0]})")
     53 self.tokenizer = tokenizer
     54 self.unk_token_id: int | None

ValueError: Number of tokens (249999) does not match number of vectors (250000)

It also doesn't work when I check beforehand whether the tokens are present in the vocabulary of the corresponding tokenizer and skip them. Even then, there are examples that don't work. We have +1000 brand names that we want to consider. Consequently, it's difficult for me to provide more explicit examples.