feat: add util for tokenizer pad id (#310)

* add tests for new util

* update test for coverage

* test pad_token passing

Author: stephantul

model2vec/train/base.py

@@ -11,6 +11,7 @@
 from torch.utils.data import DataLoader, Dataset
 
 from model2vec import StaticModel
+from model2vec.train.utils import get_probable_pad_token_id
 
 logger = logging.getLogger(__name__)
 
@@ -82,7 +83,7 @@ def from_pretrained(
 
     @classmethod
     def from_static_model(
-        cls: type[ModelType], *, model: StaticModel, out_dim: int = 2, pad_token: str = "[PAD]", **kwargs: Any
+        cls: type[ModelType], *, model: StaticModel, out_dim: int = 2, pad_token: str | None = None, **kwargs: Any
     ) -> ModelType:
         """Load the model from a static model."""
         model.embedding = np.nan_to_num(model.embedding)
@@ -92,9 +93,13 @@ def from_static_model(
             token_mapping = model.token_mapping.tolist()
         else:
             token_mapping = None
+        if pad_token is not None:
+            pad_id = model.tokenizer.get_vocab()[pad_token]
+        else:
+            pad_id = get_probable_pad_token_id(model.tokenizer)
         return cls(
             vectors=embeddings_converted,
-            pad_id=model.tokenizer.token_to_id(pad_token),
+            pad_id=pad_id,
             out_dim=out_dim,
             tokenizer=model.tokenizer,
             token_mapping=token_mapping,

model2vec/train/utils.py

@@ -0,0 +1,21 @@
+import logging
+
+from tokenizers import Tokenizer
+
+logger = logging.getLogger(__name__)
+
+_KNOWN_PAD_TOKENS = ("[PAD]", "<pad>")
+
+
+def get_probable_pad_token_id(tokenizer: Tokenizer) -> int:
+    """Get a probable pad token by using the padding module and falling back to guessing."""
+    if tokenizer.padding is not None:
+        return tokenizer.padding["pad_id"]
+    vocab = tokenizer.get_vocab()
+    for token in _KNOWN_PAD_TOKENS:
+        token_id = vocab.get(token)
+        if token_id is not None:
+            return token_id
+
+    logger.warning("No known pad token found, using 0 as default")
+    return 0

pyproject.toml

@@ -60,7 +60,7 @@ dev = [
     "ruff",
 ]
 
-distill = ["torch", "transformers", "scikit-learn", "skeletoken>=0.3.1"]
+distill = ["torch", "transformers", "scikit-learn", "skeletoken>=0.3.2"]
 onnx = ["onnx", "torch"]
 # train also installs inference
 train = ["torch", "lightning", "scikit-learn", "skops"]

tests/test_trainable.py

@@ -1,14 +1,17 @@
+import logging
 from tempfile import TemporaryDirectory
 
 import numpy as np
 import pytest
 import torch
+from skeletoken import TokenizerModel
 from tokenizers import Tokenizer
 from transformers import AutoTokenizer
 
 from model2vec.model import StaticModel
 from model2vec.train import StaticModelForClassification
 from model2vec.train.base import FinetunableStaticModel, TextDataset
+from model2vec.train.utils import get_probable_pad_token_id
 
 
 @pytest.mark.parametrize("n_layers", [0, 1, 2, 3])
@@ -67,6 +70,21 @@ def test_init_classifier_from_model(mock_vectors: np.ndarray, mock_tokenizer: To
         assert s.w.shape[0] == mock_vectors.shape[0]
 
 
+def test_pad_token(mock_tokenizer: Tokenizer) -> None:
+    """Test initializion from a static model."""
+    tokenizer_model = TokenizerModel.from_tokenizer(mock_tokenizer)
+    tokenizer_model.pad_token = "[HELLO]"
+    tokenizer = tokenizer_model.to_tokenizer()
+    vectors = np.random.RandomState().randn(6, 10)
+    model = StaticModel(vectors=vectors, tokenizer=tokenizer)
+    s = StaticModelForClassification.from_static_model(model=model, pad_token="[HELLO]")
+    assert s.w.shape[0] == vectors.shape[0]
+    assert s.pad_id == 5
+
+    with pytest.raises(KeyError):
+        StaticModelForClassification.from_static_model(model=model, pad_token="[BRR]")
+
+
 def test_encode(mock_trained_pipeline: StaticModelForClassification) -> None:
     """Test the encode function."""
     result = mock_trained_pipeline._encode(torch.tensor([[0, 1], [1, 0]]).long())
@@ -231,3 +249,35 @@ def test_evaluate(mock_trained_pipeline: StaticModelForClassification) -> None:
         else:
             # Ignore the type error since we don't support int labels in our typing, but the code does
             mock_trained_pipeline.evaluate(["dog cat", "dog"], [1, 1])  # type: ignore
+
+
+def test_get_probable_pad_token_id(mock_tokenizer: Tokenizer, caplog: pytest.LogCaptureFixture) -> None:
+    """Test loading from a static model with a pad token."""
+    tokenizer_model = TokenizerModel.from_tokenizer(mock_tokenizer)
+    t = tokenizer_model.to_tokenizer()
+    token_id = get_probable_pad_token_id(t)
+    assert token_id == 0
+
+    # Adds new token
+    tokenizer_model.pad_token = "haha"
+    t = tokenizer_model.to_tokenizer()
+    token_id = get_probable_pad_token_id(t)
+    assert token_id == 5
+
+    tokenizer_model.pad_token = "word1"
+    t = tokenizer_model.to_tokenizer()
+    token_id = get_probable_pad_token_id(t)
+    assert token_id == 1
+
+    # Remove padding token
+    tokenizer_model.pad_token = None
+    t = tokenizer_model.to_tokenizer()
+    token_id = get_probable_pad_token_id(t)
+    assert token_id == tokenizer_model.vocabulary["[PAD]"]
+
+    tokenizer_model = tokenizer_model.remove_token_from_vocabulary("[PAD]")
+    t = tokenizer_model.to_tokenizer()
+    with caplog.at_level(logging.WARNING, logger="model2vec.train.utils"):
+        token_id = get_probable_pad_token_id(t)
+    assert token_id == 0
+    assert "No known pad token found, using 0 as default" in caplog.text