feat: Integration or regression tests

.github/workflows/ci.yaml

@@ -50,7 +50,7 @@ jobs:
       - name: Run tests under coverage
         shell: bash
         run: |
-          coverage run --source=model2vec -m pytest
+          coverage run --source=model2vec -m pytest --ignore=tests/integration
           coverage report
 
       - name: Upload results to Codecov

Makefile

@@ -8,7 +8,7 @@ install:
 	uv run pre-commit install
 
 install-no-pre-commit:
-	uv pip install ".[dev,distill,inference,train,onnx,quantization]"
+	uv pip install ".[dev,distill,inference,train,onnx,quantization,integration]"
 
 install-base:
 	uv sync --extra dev
@@ -17,7 +17,13 @@ fix:
 	uv run pre-commit run --all-files
 
 test:
-	uv run pytest --cov=model2vec --cov-report=term-missing $(VERBOSITY)
+	uv run pytest --cov=model2vec --cov-report=term-missing --ignore=tests/integration $(VERBOSITY)
 
 test-verbose:
 	make test VERBOSITY="-vvv"
+
+test-integration:
+	uv run pytest tests/integration $(VERBOSITY)
+
+test-integration-update:
+	uv run python -m tests.integration.update_distill_baseline

model2vec/tokenizer/tokenizer.py

@@ -41,7 +41,7 @@ def clean_and_create_vocabulary(
     tokens_to_add: list[str] = []
     added_tokens_to_add: list[str] = []
     for token in vocabulary_to_add:
-        preprocessed = preprocessor.preprocess(token)
+        preprocessed = preprocessor.preprocess(token, had_word_prefix=True)
         if len(preprocessed) < 1:
             logger.warning(f"Token '{token}' was empty after preprocessing.")
             n_empty += 1
@@ -69,9 +69,13 @@ def clean_and_create_vocabulary(
         seen_tokens.add(token)
         tokens_to_add.append(token)
 
+    # Remove the post processor.
     model.post_processor = None
+    # Remove all prior added tokens
     model = model.prune_added_tokens()
-    model = model.add_tokens_to_vocabulary(tokens_to_add, preprocess_tokens=True)
+    # Preprocess tokens is False because tokens are already preprocessed.
+    model = model.add_tokens_to_vocabulary(tokens_to_add, preprocess_tokens=False)
+    # The tokens are not special tokens, not single words, and already normalized.
     model = model.add_addedtokens(added_tokens_to_add, is_special=False, single_word=False, normalized=True)
 
     n_multiword = len(added_tokens_to_add)

pyproject.toml

@@ -58,12 +58,13 @@ dev = [
     "setuptools",
 ]
 
-distill = ["torch", "transformers<5.4.0", "scikit-learn", "skeletoken>=0.3.3"]
+distill = ["torch", "transformers<5.4.0", "scikit-learn", "skeletoken>=0.4.0,<0.5.0"]
 onnx = ["onnx", "torch"]
 # train also installs inference
 train = ["torch", "lightning", "scikit-learn", "skops"]
 inference = ["scikit-learn", "skops"]
 quantization = ["scikit-learn"]
+integration = ["mteb"]
 
 [project.urls]
 "Homepage" = "https://github.com/MinishLab"
@@ -123,3 +124,4 @@ version = {attr = "model2vec.version.__version__"}
 
 [tool.uv]
 exclude-newer = "1 week"
+exclude-newer-package = { skeletoken = false }

tests/conftest.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import Any, cast
+from typing import Any, Iterator, cast
 
 import numpy as np
 import pytest
@@ -62,6 +62,14 @@ def mock_tokenizermodel() -> TokenizerModel:
     return TokenizerModel.from_pretrained("tests/data/test_tokenizer")
 
 
+# TODO: Temporary fix for skeletoken 0.3.3 compatibility.
+# Once we go to >= 0.4.0 this is no longer needed
+class DumbConfig:
+    def __iter__(self) -> Iterator[tuple[str, object]]:
+        """Yield no config values."""
+        yield from []
+
+
 @pytest.fixture
 def mock_transformer(request: pytest.FixtureRequest) -> PreTrainedModel:
     """Create a mock transformer model."""
@@ -80,7 +88,7 @@ def __init__(self, vocab_size: int, dim: int, with_pooler: bool, pooler_value: f
             self.with_pooler = with_pooler
             self.pooler_value = pooler_value
             self.input_embs = torch.nn.Embedding(vocab_size, dim)
-            self.config: dict[str, Any] = {}
+            self.config = DumbConfig()
 
         def to(self, device: str) -> MockPreTrainedModel:
             self.device = device

tests/integration/data/Alibaba-NLP___gte-modernbert-base_baseline.json

@@ -0,0 +1,91 @@
+{
+  "base_model": "Alibaba-NLP/gte-modernbert-base",
+  "configs": {
+    "custom_vocab": {
+      "embedding_dim": 32,
+      "embedding_mean": 0.000534,
+      "embedding_rank": 32,
+      "embedding_rows": 50269,
+      "embedding_std": 3.379484,
+      "first_tokens": [
+        "!",
+        "\"",
+        "#",
+        "$",
+        "%",
+        "&",
+        "'",
+        "(",
+        ")",
+        "*"
+      ],
+      "full_vocab_size": 50269,
+      "last_tokens": [
+        "\u0120quixnorfle",
+        "norfle",
+        "\u0120quix",
+        "\u0120blorptastic",
+        "ptastic",
+        "\u0120blor",
+        "\u0120flimzycrag",
+        "zycr",
+        "\u0120flim",
+        "\u0120flimzycr"
+      ],
+      "mteb_sts_scores": {
+        "BIOSSES": 0.369924,
+        "SICK-R": 0.496052,
+        "STS12": 0.499049,
+        "STS13": 0.586944,
+        "STS14": 0.526949,
+        "STS15": 0.628959,
+        "STS16": 0.596379,
+        "STSBenchmark": 0.536571
+      },
+      "token_order_hash": "c2f9a27e98ef519d36c74cbde4ccc7ab023a23896aac6351f30c41c900daac78"
+    },
+    "subword": {
+      "embedding_dim": 256,
+      "embedding_mean": 3.1e-05,
+      "embedding_rank": 256,
+      "embedding_rows": 50254,
+      "embedding_std": 1.690459,
+      "first_tokens": [
+        "!",
+        "\"",
+        "#",
+        "$",
+        "%",
+        "&",
+        "'",
+        "(",
+        ")",
+        "*"
+      ],
+      "full_vocab_size": 50254,
+      "last_tokens": [
+        "767",
+        "itons",
+        "\u0120PIP",
+        "\u0120Tus",
+        "ibrated",
+        "\u0120fortified",
+        "ferenced",
+        "\u0120Outcomes",
+        "[UNK]",
+        "[PAD]"
+      ],
+      "mteb_sts_scores": {
+        "BIOSSES": 0.525464,
+        "SICK-R": 0.553442,
+        "STS12": 0.547261,
+        "STS13": 0.659771,
+        "STS14": 0.615254,
+        "STS15": 0.708012,
+        "STS16": 0.669369,
+        "STSBenchmark": 0.609307
+      },
+      "token_order_hash": "c3ba2855d4683a2c49936b821bd0d7ad2f137bf090f66ee382bb85c350b4baba"
+    }
+  }
+}

tests/integration/data/baai___bge-base-en-v1.5_baseline.json

@@ -0,0 +1,91 @@
+{
+  "base_model": "baai/bge-base-en-v1.5",
+  "configs": {
+    "custom_vocab": {
+      "embedding_dim": 32,
+      "embedding_mean": -0.00102,
+      "embedding_rank": 32,
+      "embedding_rows": 29529,
+      "embedding_std": 0.907232,
+      "first_tokens": [
+        "[PAD]",
+        "[UNK]",
+        "!",
+        "\"",
+        "#",
+        "$",
+        "%",
+        "&",
+        "'",
+        "("
+      ],
+      "full_vocab_size": 29529,
+      "last_tokens": [
+        "##\uff0d",
+        "##\uff0e",
+        "##\uff0f",
+        "##\uff1a",
+        "##\uff1f",
+        "##\uff5e",
+        "zibblorptron",
+        "quixnorfle",
+        "blorptastic",
+        "flimzycrag"
+      ],
+      "mteb_sts_scores": {
+        "BIOSSES": 0.60555,
+        "SICK-R": 0.574381,
+        "STS12": 0.615507,
+        "STS13": 0.715354,
+        "STS14": 0.638972,
+        "STS15": 0.682774,
+        "STS16": 0.623959,
+        "STSBenchmark": 0.611345
+      },
+      "token_order_hash": "09d9aaaedc7d41b2fd1e3ffa64f16794a2cefe9ffc7faf1055ae5d49350c3060"
+    },
+    "subword": {
+      "embedding_dim": 256,
+      "embedding_mean": -0.000118,
+      "embedding_rank": 256,
+      "embedding_rows": 29525,
+      "embedding_std": 0.562719,
+      "first_tokens": [
+        "[PAD]",
+        "[UNK]",
+        "!",
+        "\"",
+        "#",
+        "$",
+        "%",
+        "&",
+        "'",
+        "("
+      ],
+      "full_vocab_size": 29525,
+      "last_tokens": [
+        "##\uff01",
+        "##\uff08",
+        "##\uff09",
+        "##\uff0c",
+        "##\uff0d",
+        "##\uff0e",
+        "##\uff0f",
+        "##\uff1a",
+        "##\uff1f",
+        "##\uff5e"
+      ],
+      "mteb_sts_scores": {
+        "BIOSSES": 0.742493,
+        "SICK-R": 0.633061,
+        "STS12": 0.634404,
+        "STS13": 0.755101,
+        "STS14": 0.697645,
+        "STS15": 0.777181,
+        "STS16": 0.708249,
+        "STSBenchmark": 0.694604
+      },
+      "token_order_hash": "d881fab650dd73240615f87f531ff31a6135f961933ebc32c0dfdd130763e7a8"
+    }
+  }
+}

tests/integration/data/google___embeddinggemma-300m_baseline.json

@@ -0,0 +1,91 @@
+{
+  "base_model": "google/embeddinggemma-300m",
+  "configs": {
+    "custom_vocab": {
+      "embedding_dim": 32,
+      "embedding_mean": 0.000445,
+      "embedding_rank": 32,
+      "embedding_rows": 255746,
+      "embedding_std": 2.829036,
+      "first_tokens": [
+        "<pad>",
+        "<unk>",
+        "<0x00>",
+        "<0x01>",
+        "<0x02>",
+        "<0x03>",
+        "<0x04>",
+        "<0x05>",
+        "<0x06>",
+        "<0x07>"
+      ],
+      "full_vocab_size": 255746,
+      "last_tokens": [
+        "zibblor",
+        "quixnorfle",
+        "norfle",
+        "quix",
+        "blorptastic",
+        "blorpt",
+        "flimzycrag",
+        "flim",
+        "flimzycr",
+        "zycr"
+      ],
+      "mteb_sts_scores": {
+        "BIOSSES": 0.48233,
+        "SICK-R": 0.570594,
+        "STS12": 0.519418,
+        "STS13": 0.622716,
+        "STS14": 0.537778,
+        "STS15": 0.575234,
+        "STS16": 0.534645,
+        "STSBenchmark": 0.451954
+      },
+      "token_order_hash": "bb8f48741104cf94146b583abb11a896d6005dbf7e04403053821a12ddeecf42"
+    },
+    "subword": {
+      "embedding_dim": 256,
+      "embedding_mean": 1.4e-05,
+      "embedding_rank": 256,
+      "embedding_rows": 255732,
+      "embedding_std": 1.653758,
+      "first_tokens": [
+        "<pad>",
+        "<unk>",
+        "<0x00>",
+        "<0x01>",
+        "<0x02>",
+        "<0x03>",
+        "<0x04>",
+        "<0x05>",
+        "<0x06>",
+        "<0x07>"
+      ],
+      "full_vocab_size": 255732,
+      "last_tokens": [
+        "\ud83e\udd3d",
+        "\u13dc",
+        "\u141b",
+        "\u257e",
+        "\u25a5",
+        "\u7c59",
+        "\u8182",
+        "\u87f3",
+        "\ud44e",
+        "\ud835\udd39"
+      ],
+      "mteb_sts_scores": {
+        "BIOSSES": 0.639707,
+        "SICK-R": 0.633239,
+        "STS12": 0.566534,
+        "STS13": 0.731087,
+        "STS14": 0.647276,
+        "STS15": 0.723692,
+        "STS16": 0.644938,
+        "STSBenchmark": 0.58957
+      },
+      "token_order_hash": "cf0c61603b3c2786463e128bb975616001935355803924492e245746455ed235"
+    }
+  }
+}