diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index 5ded6cd..cf9224f 100644 --- a/.github/workflows/ci.yaml +++ b/.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 diff --git a/Makefile b/Makefile index 276817e..c067b77 100644 --- a/Makefile +++ b/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 diff --git a/model2vec/tokenizer/tokenizer.py b/model2vec/tokenizer/tokenizer.py index c939a3b..463f171 100644 --- a/model2vec/tokenizer/tokenizer.py +++ b/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) diff --git a/pyproject.toml b/pyproject.toml index dd8cf83..0a9e75e 100644 --- a/pyproject.toml +++ b/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 } diff --git a/tests/conftest.py b/tests/conftest.py index 2f45eef..af46343 100644 --- a/tests/conftest.py +++ b/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 diff --git a/tests/integration/__init__.py b/tests/integration/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/integration/data/Alibaba-NLP___gte-modernbert-base_baseline.json b/tests/integration/data/Alibaba-NLP___gte-modernbert-base_baseline.json new file mode 100644 index 0000000..0f10306 --- /dev/null +++ b/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" + } + } +} diff --git a/tests/integration/data/baai___bge-base-en-v1.5_baseline.json b/tests/integration/data/baai___bge-base-en-v1.5_baseline.json new file mode 100644 index 0000000..da97ebf --- /dev/null +++ b/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" + } + } +} diff --git a/tests/integration/data/google___embeddinggemma-300m_baseline.json b/tests/integration/data/google___embeddinggemma-300m_baseline.json new file mode 100644 index 0000000..e1003ec --- /dev/null +++ b/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": [ + "", + "", + "<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": [ + "", + "", + "<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" + } + } +} diff --git a/tests/integration/data/intfloat___multilingual-e5-base_baseline.json b/tests/integration/data/intfloat___multilingual-e5-base_baseline.json new file mode 100644 index 0000000..d2299a5 --- /dev/null +++ b/tests/integration/data/intfloat___multilingual-e5-base_baseline.json @@ -0,0 +1,91 @@ +{ + "base_model": "intfloat/multilingual-e5-base", + "configs": { + "custom_vocab": { + "embedding_dim": 32, + "embedding_mean": -0.00019, + "embedding_rank": 32, + "embedding_rows": 250003, + "embedding_std": 0.764859, + "first_tokens": [ + "", + "", + ",", + ".", + "\u2581", + "s", + "\u2581de", + "-", + "\u2581a", + "a" + ], + "full_vocab_size": 250003, + "last_tokens": [ + "\u5a20", + "\u5f8a", + "\u8718", + "\u8e34", + "\u1ea5", + "\u7a23", + "\u2581zibblorptron", + "\u2581quixnorfle", + "\u2581blorptastic", + "\u2581flimzycrag" + ], + "mteb_sts_scores": { + "BIOSSES": 0.432685, + "SICK-R": 0.474725, + "STS12": 0.520844, + "STS13": 0.623774, + "STS14": 0.565683, + "STS15": 0.635075, + "STS16": 0.58762, + "STSBenchmark": 0.492349 + }, + "token_order_hash": "c15b66edc095c19fae1fc0d10462ccd3cdc77e835825639022dac2e8a3f4b5bd" + }, + "subword": { + "embedding_dim": 256, + "embedding_mean": 1.1e-05, + "embedding_rank": 256, + "embedding_rows": 249999, + "embedding_std": 0.417482, + "first_tokens": [ + "", + "", + ",", + ".", + "\u2581", + "s", + "\u2581de", + "-", + "\u2581a", + "a" + ], + "full_vocab_size": 249999, + "last_tokens": [ + "\u5783", + "\u5f98", + "\u8d41", + "\u1ea7", + "\u5a20", + "\u5f8a", + "\u8718", + "\u8e34", + "\u1ea5", + "\u7a23" + ], + "mteb_sts_scores": { + "BIOSSES": 0.512096, + "SICK-R": 0.556843, + "STS12": 0.583896, + "STS13": 0.706956, + "STS14": 0.664318, + "STS15": 0.735835, + "STS16": 0.682632, + "STSBenchmark": 0.625896 + }, + "token_order_hash": "d0534dadde49bec613f7090e0fec5797f10194566612acf36895214252e290e0" + } + } +} diff --git a/tests/integration/data/sentence-transformers___all-MiniLM-L6-v2_baseline.json b/tests/integration/data/sentence-transformers___all-MiniLM-L6-v2_baseline.json new file mode 100644 index 0000000..50a54a0 --- /dev/null +++ b/tests/integration/data/sentence-transformers___all-MiniLM-L6-v2_baseline.json @@ -0,0 +1,91 @@ +{ + "base_model": "sentence-transformers/all-MiniLM-L6-v2", + "configs": { + "custom_vocab": { + "embedding_dim": 32, + "embedding_mean": -0.000871, + "embedding_rank": 32, + "embedding_rows": 29529, + "embedding_std": 0.63969, + "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.500132, + "SICK-R": 0.565143, + "STS12": 0.594862, + "STS13": 0.653794, + "STS14": 0.596518, + "STS15": 0.611507, + "STS16": 0.569707, + "STSBenchmark": 0.559639 + }, + "token_order_hash": "09d9aaaedc7d41b2fd1e3ffa64f16794a2cefe9ffc7faf1055ae5d49350c3060" + }, + "subword": { + "embedding_dim": 256, + "embedding_mean": -0.0001, + "embedding_rank": 256, + "embedding_rows": 29525, + "embedding_std": 0.368635, + "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.656286, + "SICK-R": 0.611602, + "STS12": 0.620365, + "STS13": 0.73547, + "STS14": 0.670425, + "STS15": 0.745883, + "STS16": 0.67784, + "STSBenchmark": 0.655636 + }, + "token_order_hash": "d881fab650dd73240615f87f531ff31a6135f961933ebc32c0dfdd130763e7a8" + } + } +} diff --git a/tests/integration/distill_metrics.py b/tests/integration/distill_metrics.py new file mode 100644 index 0000000..7b8ee4c --- /dev/null +++ b/tests/integration/distill_metrics.py @@ -0,0 +1,123 @@ +from __future__ import annotations + +import hashlib +from pathlib import Path +from typing import Any, cast + +import mteb +import numpy as np +from mteb.models.abs_encoder import AbsEncoder +from mteb.models.model_meta import ModelMeta +from transformers import AutoModel, AutoTokenizer +from transformers.modeling_utils import PreTrainedModel +from transformers.tokenization_utils_fast import PreTrainedTokenizerFast + +from model2vec.distill import distill_from_model +from model2vec.model import StaticModel + +BASE_MODELS: tuple[str, ...] = ( + "sentence-transformers/all-MiniLM-L6-v2", + "baai/bge-base-en-v1.5", + "intfloat/multilingual-e5-base", + "google/embeddinggemma-300m", + "Alibaba-NLP/gte-modernbert-base", +) + +BASELINE_DIR = Path(__file__).parent / "data" + +STS_TASKS: tuple[str, ...] = ( + "BIOSSES", + "SICK-R", + "STS12", + "STS13", + "STS14", + "STS15", + "STS16", + "STSBenchmark", +) + +_NOVEL_VOCABULARY = ["zibblorptron", "quixnorfle", "blorptastic", "flimzycrag"] + +CONFIGS: dict[str, dict[str, Any]] = { + "subword": {"pca_dims": 256, "quantize_to": "float32"}, + "custom_vocab": {"vocabulary": _NOVEL_VOCABULARY, "pca_dims": 32, "quantize_to": "float32"}, +} + + +def baseline_path_for(model_name: str) -> Path: + """The JSON baseline file for a given short model name (a key of `BASE_MODELS`).""" + safe_model_name = model_name.replace("/", "___") + return BASELINE_DIR / f"{safe_model_name}_baseline.json" + + +def load_base_model_and_tokenizer(model_name: str) -> tuple[PreTrainedModel, PreTrainedTokenizerFast]: + """Download a base sentence-transformer and its tokenizer once, for reuse across distillations. + + :param model_name: The HuggingFace model id to download, e.g. `BASE_MODELS["minilm"]`. + :return: The loaded model and tokenizer. + """ + model = AutoModel.from_pretrained(model_name) + tokenizer = cast(PreTrainedTokenizerFast, AutoTokenizer.from_pretrained(model_name, use_fast=True)) + return model, tokenizer + + +def distill_all(model: PreTrainedModel, tokenizer: PreTrainedTokenizerFast) -> dict[str, StaticModel]: + """Distill every configured variant from the same base model and tokenizer.""" + return {name: distill_from_model(model=model, tokenizer=tokenizer, **kwargs) for name, kwargs in CONFIGS.items()} + + +class _MTEBEncoder(AbsEncoder): + """Adapts a `StaticModel` to MTEB's `AbsEncoder` protocol.""" + + def __init__(self, model: StaticModel) -> None: + self.model = model + self.mteb_model_meta = ModelMeta.create_empty(overwrites={"name": "model2vec-distilled", "revision": "local"}) + + def encode( + self, + inputs: Any, + *, + task_metadata: Any, + hf_split: str, + hf_subset: str, + prompt_type: Any = None, + **kwargs: Any, + ) -> np.ndarray: + texts = [text for batch in inputs for text in batch["text"]] + return self.model.encode(texts) + + +def compute_mteb_sts_scores(model: StaticModel) -> dict[str, float]: + """Run the STS subset of MTEB against a distilled model. + + :param model: The distilled StaticModel to evaluate. + :return: A dict mapping MTEB STS task name to its main score. + """ + tasks = mteb.get_tasks(tasks=list(STS_TASKS)) + results = mteb.evaluate(_MTEBEncoder(model), tasks, cache=None, show_progress_bar=False) # type: ignore[arg-type] + return {result.task_name: round(float(result.get_score()), 6) for result in results.task_results} + + +def compute_metrics(model: StaticModel) -> dict[str, Any]: + """Compute a JSON-serializable snapshot of a distilled model's key properties. + + :param model: The distilled StaticModel to summarize. + :return: A dict with vocab size, embedding shape/rank/distribution, token order, and MTEB STS + scores. Used both to write and to check the regression baseline. + """ + embedding = model.embedding.astype(np.float64) + tokens = list(model.tokens) + token_order_hash = hashlib.sha256("\x1f".join(tokens).encode("utf-8")).hexdigest() + + return { + "full_vocab_size": len(tokens), + "embedding_rows": int(embedding.shape[0]), + "embedding_dim": int(embedding.shape[1]), + "embedding_rank": int(np.linalg.matrix_rank(embedding)), + "embedding_mean": round(float(embedding.mean()), 6), + "embedding_std": round(float(embedding.std()), 6), + "token_order_hash": token_order_hash, + "first_tokens": tokens[:10], + "last_tokens": tokens[-10:], + "mteb_sts_scores": compute_mteb_sts_scores(model), + } diff --git a/tests/integration/test_distill_regression.py b/tests/integration/test_distill_regression.py new file mode 100644 index 0000000..6d178d5 --- /dev/null +++ b/tests/integration/test_distill_regression.py @@ -0,0 +1,145 @@ +from __future__ import annotations + +import copy +import json +from pathlib import Path +from typing import Any + +import numpy as np +import pytest +from transformers.modeling_utils import PreTrainedModel +from transformers.tokenization_utils_fast import PreTrainedTokenizerFast + +from model2vec.distill import distill_from_model +from model2vec.model import StaticModel +from tests.integration.distill_metrics import ( + BASE_MODELS, + CONFIGS, + baseline_path_for, + compute_metrics, + distill_all, + load_base_model_and_tokenizer, +) + +_EXACT_FIELDS = ( + "full_vocab_size", + "embedding_rows", + "embedding_dim", + "token_order_hash", + "first_tokens", + "last_tokens", +) + + +@pytest.fixture(scope="module", params=sorted(BASE_MODELS), ids=sorted(BASE_MODELS)) +def model_name(request: pytest.FixtureRequest) -> str: + """A short base-model name (a key of `BASE_MODELS`), parametrizing the whole module.""" + return request.param + + +@pytest.fixture(scope="module") +def baseline(model_name: str) -> dict[str, Any]: + """Load the stored golden baseline for this base model.""" + path = baseline_path_for(model_name) + if not path.exists(): + pytest.fail(f"No baseline found at {path}. Generate one with `make test-integration-update`.") + return json.loads(path.read_text()) + + +@pytest.fixture(scope="module") +def base_model_and_tokenizer(model_name: str) -> tuple[PreTrainedModel, PreTrainedTokenizerFast]: + """Download this module's base sentence-transformer once, for reuse across every test in it.""" + return load_base_model_and_tokenizer(model_name) + + +@pytest.fixture(scope="module") +def distilled_models( + base_model_and_tokenizer: tuple[PreTrainedModel, PreTrainedTokenizerFast], +) -> dict[str, StaticModel]: + """Distill every configured variant from the base model, once for the whole module.""" + model, tokenizer = base_model_and_tokenizer + return distill_all(model, tokenizer) + + +@pytest.fixture(scope="module") +def current_metrics(distilled_models: dict[str, StaticModel]) -> dict[str, dict[str, Any]]: + """Compute the same metrics as the baseline for the freshly distilled models.""" + return {name: compute_metrics(static_model) for name, static_model in distilled_models.items()} + + +@pytest.mark.parametrize("config_name", sorted(CONFIGS)) +def test_structural_metrics_match_baseline( + config_name: str, model_name: str, baseline: dict[str, Any], current_metrics: dict[str, dict[str, Any]] +) -> None: + """Vocab size, embedding shape, token order, and token identity must match the baseline exactly.""" + expected = baseline["configs"][config_name] + actual = current_metrics[config_name] + for field in _EXACT_FIELDS: + assert actual[field] == expected[field], ( + f"[{model_name}/{config_name}] '{field}' drifted from the baseline: " + f"expected {expected[field]!r}, got {actual[field]!r}. " + "If this is intentional, run `make test-integration-update` and review the JSON diff." + ) + + +@pytest.mark.parametrize("config_name", sorted(CONFIGS)) +def test_embedding_rank_matches_baseline( + config_name: str, model_name: str, baseline: dict[str, Any], current_metrics: dict[str, dict[str, Any]] +) -> None: + """A rank drop vs. the baseline means the pooling/PCA/quantization step degenerated.""" + expected_rank = baseline["configs"][config_name]["embedding_rank"] + actual_rank = current_metrics[config_name]["embedding_rank"] + assert abs(actual_rank - expected_rank) <= 1, ( + f"[{model_name}/{config_name}] embedding rank {actual_rank} vs baseline {expected_rank}" + ) + + +@pytest.mark.parametrize("config_name", sorted(CONFIGS)) +def test_embedding_distribution_matches_baseline( + config_name: str, model_name: str, baseline: dict[str, Any], current_metrics: dict[str, dict[str, Any]] +) -> None: + """The embedding matrix's mean/std shouldn't drift far from the baseline (small tolerance for numerical noise).""" + expected = baseline["configs"][config_name] + actual = current_metrics[config_name] + assert actual["embedding_mean"] == pytest.approx(expected["embedding_mean"], abs=1e-3), model_name + assert actual["embedding_std"] == pytest.approx(expected["embedding_std"], rel=0.05, abs=1e-3), model_name + + +@pytest.mark.parametrize("config_name", sorted(CONFIGS)) +def test_mteb_sts_scores_match_baseline( + config_name: str, model_name: str, baseline: dict[str, Any], current_metrics: dict[str, dict[str, Any]] +) -> None: + """Repeated distillation must score the same on the MTEB STS tasks as the golden baseline.""" + expected_scores = baseline["configs"][config_name]["mteb_sts_scores"] + actual_scores = current_metrics[config_name]["mteb_sts_scores"] + for task_name, expected_score in expected_scores.items(): + actual_score = actual_scores[task_name] + assert actual_score == pytest.approx(expected_score, abs=0.01), ( + f"[{model_name}/{config_name}] MTEB '{task_name}' STS score drifted from baseline: " + f"{expected_score:.4f} -> {actual_score:.4f}" + ) + + +def test_distillation_is_deterministic( + base_model_and_tokenizer: tuple[PreTrainedModel, PreTrainedTokenizerFast], +) -> None: + """Distilling twice with identical parameters must yield identical embeddings and vocabularies.""" + model, tokenizer = base_model_and_tokenizer + kwargs = CONFIGS["custom_vocab"] + + first = distill_from_model(model=copy.deepcopy(model), tokenizer=tokenizer, **kwargs) + second = distill_from_model(model=copy.deepcopy(model), tokenizer=tokenizer, **kwargs) + + assert first.tokens == second.tokens + np.testing.assert_allclose(first.embedding, second.embedding, rtol=1e-5, atol=1e-6) + + +def test_save_and_load_roundtrip(distilled_models: dict[str, StaticModel], tmp_path: Path) -> None: + """Saving and reloading a distilled model must not change its tokens or embeddings.""" + model = distilled_models["subword"] + save_path = tmp_path / "distilled_model" + model.save_pretrained(save_path) + loaded_model = StaticModel.from_pretrained(save_path) + + assert loaded_model.tokens == model.tokens + np.testing.assert_array_equal(loaded_model.embedding, model.embedding) diff --git a/tests/integration/update_distill_baseline.py b/tests/integration/update_distill_baseline.py new file mode 100644 index 0000000..8a0dee9 --- /dev/null +++ b/tests/integration/update_distill_baseline.py @@ -0,0 +1,47 @@ +from __future__ import annotations + +import json +import logging +import sys + +from tests.integration.distill_metrics import ( + BASE_MODELS, + baseline_path_for, + compute_metrics, + distill_all, + load_base_model_and_tokenizer, +) + +logger = logging.getLogger(__name__) + + +def update_baseline(model_name: str) -> None: + """Distill every configured variant of one base model and write their metrics to its baseline file. + + :param model_name: A model identifier on Hugging Face. + """ + model, tokenizer = load_base_model_and_tokenizer(model_name) + distilled = distill_all(model, tokenizer) + + baseline = { + "base_model": model_name, + "configs": {name: compute_metrics(static_model) for name, static_model in distilled.items()}, + } + + path = baseline_path_for(model_name) + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(baseline, indent=2, sort_keys=True) + "\n") + logger.info(f"Wrote baseline for '{model_name}' to {path}") + + +def main() -> None: + """Regenerate the baselines requested on the command line, or all of them if none were given.""" + requested = sys.argv[1:] or sorted(BASE_MODELS) + + for model_name in requested: + update_baseline(model_name) + + +if __name__ == "__main__": + logging.basicConfig(level=logging.INFO, format="%(message)s") + main() diff --git a/uv.lock b/uv.lock index 000d536..28f802b 100644 --- a/uv.lock +++ b/uv.lock @@ -12,6 +12,9 @@ resolution-markers = [ exclude-newer = "0001-01-01T00:00:00Z" # This has no effect and is included for backwards compatibility when using relative exclude-newer values. exclude-newer-span = "P3D" +[options.exclude-newer-package] +skeletoken = false + [[package]] name = "aiohappyeyeballs" version = "2.6.1" @@ -418,7 +421,7 @@ name = "exceptiongroup" version = "1.3.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "typing-extensions", marker = "python_full_version < '3.11'" }, + { name = "typing-extensions" }, ] sdist = { url = "https://files.pythonhosted.org/packages/50/79/66800aadf48771f6b62f7eb014e352e5d06856655206165d775e675a02c9/exceptiongroup-1.3.1.tar.gz", hash = "sha256:8b412432c6055b0b7d14c310000ae93352ed6754f70fa8f7c34141f91c4e3219", size = 30371, upload-time = "2025-11-21T23:01:54.787Z" } wheels = [ @@ -702,17 +705,17 @@ resolution-markers = [ "python_full_version < '3.11'", ] dependencies = [ - { name = "colorama", marker = "python_full_version < '3.11' and sys_platform == 'win32'" }, - { name = "decorator", marker = "python_full_version < '3.11'" }, - { name = "exceptiongroup", marker = "python_full_version < '3.11'" }, - { name = "jedi", marker = "python_full_version < '3.11'" }, - { name = "matplotlib-inline", marker = "python_full_version < '3.11'" }, - { name = "pexpect", marker = "python_full_version < '3.11' and sys_platform != 'emscripten' and sys_platform != 'win32'" }, - { name = "prompt-toolkit", marker = "python_full_version < '3.11'" }, - { name = "pygments", marker = "python_full_version < '3.11'" }, - { name = "stack-data", marker = "python_full_version < '3.11'" }, - { name = "traitlets", marker = "python_full_version < '3.11'" }, - { name = "typing-extensions", marker = "python_full_version < '3.11'" }, + { name = "colorama", marker = "sys_platform == 'win32'" }, + { name = "decorator" }, + { name = "exceptiongroup" }, + { name = "jedi" }, + { name = "matplotlib-inline" }, + { name = "pexpect", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, + { name = "prompt-toolkit" }, + { name = "pygments" }, + { name = "stack-data" }, + { name = "traitlets" }, + { name = "typing-extensions" }, ] sdist = { url = "https://files.pythonhosted.org/packages/40/18/f8598d287006885e7136451fdea0755af4ebcbfe342836f24deefaed1164/ipython-8.39.0.tar.gz", hash = "sha256:4110ae96012c379b8b6db898a07e186c40a2a1ef5d57a7fa83166047d9da7624", size = 5513971, upload-time = "2026-03-27T10:02:13.94Z" } wheels = [ @@ -727,17 +730,17 @@ resolution-markers = [ "python_full_version == '3.11.*'", ] dependencies = [ - { name = "colorama", marker = "python_full_version == '3.11.*' and sys_platform == 'win32'" }, - { name = "decorator", marker = "python_full_version == '3.11.*'" }, - { name = "ipython-pygments-lexers", marker = "python_full_version == '3.11.*'" }, - { name = "jedi", marker = "python_full_version == '3.11.*'" }, - { name = "matplotlib-inline", marker = "python_full_version == '3.11.*'" }, - { name = "pexpect", marker = "python_full_version == '3.11.*' and sys_platform != 'emscripten' and sys_platform != 'win32'" }, - { name = "prompt-toolkit", marker = "python_full_version == '3.11.*'" }, - { name = "pygments", marker = "python_full_version == '3.11.*'" }, - { name = "stack-data", marker = "python_full_version == '3.11.*'" }, - { name = "traitlets", marker = "python_full_version == '3.11.*'" }, - { name = "typing-extensions", marker = "python_full_version == '3.11.*'" }, + { name = "colorama", marker = "sys_platform == 'win32'" }, + { name = "decorator" }, + { name = "ipython-pygments-lexers" }, + { name = "jedi" }, + { name = "matplotlib-inline" }, + { name = "pexpect", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, + { name = "prompt-toolkit" }, + { name = "pygments" }, + { name = "stack-data" }, + { name = "traitlets" }, + { name = "typing-extensions" }, ] sdist = { url = "https://files.pythonhosted.org/packages/c5/25/daae0e764047b0a2480c7bbb25d48f4f509b5818636562eeac145d06dfee/ipython-9.10.1.tar.gz", hash = "sha256:e170e9b2a44312484415bdb750492699bf329233b03f2557a9692cce6466ada4", size = 4426663, upload-time = "2026-03-27T09:53:26.244Z" } wheels = [ @@ -753,16 +756,16 @@ resolution-markers = [ "python_full_version == '3.12.*'", ] dependencies = [ - { name = "colorama", marker = "python_full_version >= '3.12' and sys_platform == 'win32'" }, - { name = "decorator", marker = "python_full_version >= '3.12'" }, - { name = "ipython-pygments-lexers", marker = "python_full_version >= '3.12'" }, - { name = "jedi", marker = "python_full_version >= '3.12'" }, - { name = "matplotlib-inline", marker = "python_full_version >= '3.12'" }, - { name = "pexpect", marker = "python_full_version >= '3.12' and sys_platform != 'emscripten' and sys_platform != 'win32'" }, - { name = "prompt-toolkit", marker = "python_full_version >= '3.12'" }, - { name = "pygments", marker = "python_full_version >= '3.12'" }, - { name = "stack-data", marker = "python_full_version >= '3.12'" }, - { name = "traitlets", marker = "python_full_version >= '3.12'" }, + { name = "colorama", marker = "sys_platform == 'win32'" }, + { name = "decorator" }, + { name = "ipython-pygments-lexers" }, + { name = "jedi" }, + { name = "matplotlib-inline" }, + { name = "pexpect", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, + { name = "prompt-toolkit" }, + { name = "pygments" }, + { name = "stack-data" }, + { name = "traitlets" }, ] sdist = { url = "https://files.pythonhosted.org/packages/3a/73/7114f80a8f9cabdb13c27732dce24af945b2923dcab80723602f7c8bc2d8/ipython-9.12.0.tar.gz", hash = "sha256:01daa83f504b693ba523b5a407246cabde4eb4513285a3c6acaff11a66735ee4", size = 4428879, upload-time = "2026-03-27T09:42:45.312Z" } wheels = [ @@ -774,7 +777,7 @@ name = "ipython-pygments-lexers" version = "1.1.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "pygments", marker = "python_full_version >= '3.11'" }, + { name = "pygments" }, ] sdist = { url = "https://files.pythonhosted.org/packages/ef/4c/5dd1d8af08107f88c7f741ead7a40854b8ac24ddf9ae850afbcf698aa552/ipython_pygments_lexers-1.1.1.tar.gz", hash = "sha256:09c0138009e56b6854f9535736f4171d855c8c08a563a0dcd8022f78355c7e81", size = 8393, upload-time = "2025-01-17T11:24:34.505Z" } wheels = [ @@ -1168,7 +1171,7 @@ requires-dist = [ { name = "scikit-learn", marker = "extra == 'quantization'" }, { name = "scikit-learn", marker = "extra == 'train'" }, { name = "setuptools", marker = "extra == 'dev'" }, - { name = "skeletoken", marker = "extra == 'distill'", specifier = ">=0.3.3" }, + { name = "skeletoken", marker = "extra == 'distill'", specifier = ">=0.4.0,<0.5.0" }, { name = "skops", marker = "extra == 'inference'" }, { name = "skops", marker = "extra == 'train'" }, { name = "tokenizers", specifier = ">=0.20" }, @@ -2444,10 +2447,10 @@ resolution-markers = [ "python_full_version < '3.11'", ] dependencies = [ - { name = "joblib", marker = "python_full_version < '3.11'" }, - { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, - { name = "scipy", version = "1.15.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, - { name = "threadpoolctl", marker = "python_full_version < '3.11'" }, + { name = "joblib" }, + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" } }, + { name = "scipy", version = "1.15.3", source = { registry = "https://pypi.org/simple" } }, + { name = "threadpoolctl" }, ] sdist = { url = "https://files.pythonhosted.org/packages/98/c2/a7855e41c9d285dfe86dc50b250978105dce513d6e459ea66a6aeb0e1e0c/scikit_learn-1.7.2.tar.gz", hash = "sha256:20e9e49ecd130598f1ca38a1d85090e1a600147b9c02fa6f15d69cb53d968fda", size = 7193136, upload-time = "2025-09-09T08:21:29.075Z" } wheels = [ @@ -2493,10 +2496,10 @@ resolution-markers = [ "python_full_version == '3.11.*'", ] dependencies = [ - { name = "joblib", marker = "python_full_version >= '3.11'" }, - { name = "numpy", version = "2.4.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, - { name = "scipy", version = "1.17.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, - { name = "threadpoolctl", marker = "python_full_version >= '3.11'" }, + { name = "joblib" }, + { name = "numpy", version = "2.4.3", source = { registry = "https://pypi.org/simple" } }, + { name = "scipy", version = "1.17.1", source = { registry = "https://pypi.org/simple" } }, + { name = "threadpoolctl" }, ] sdist = { url = "https://files.pythonhosted.org/packages/0e/d4/40988bf3b8e34feec1d0e6a051446b1f66225f8529b9309becaeef62b6c4/scikit_learn-1.8.0.tar.gz", hash = "sha256:9bccbb3b40e3de10351f8f5068e105d0f4083b1a65fa07b6634fbc401a6287fd", size = 7335585, upload-time = "2025-12-10T07:08:53.618Z" } wheels = [ @@ -2546,7 +2549,7 @@ resolution-markers = [ "python_full_version < '3.11'", ] dependencies = [ - { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" } }, ] sdist = { url = "https://files.pythonhosted.org/packages/0f/37/6964b830433e654ec7485e45a00fc9a27cf868d622838f6b6d9c5ec0d532/scipy-1.15.3.tar.gz", hash = "sha256:eae3cf522bc7df64b42cad3925c876e1b0b6c35c1337c93e12c0f366f55b0eaf", size = 59419214, upload-time = "2025-05-08T16:13:05.955Z" } wheels = [ @@ -2607,7 +2610,7 @@ resolution-markers = [ "python_full_version == '3.11.*'", ] dependencies = [ - { name = "numpy", version = "2.4.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "numpy", version = "2.4.3", source = { registry = "https://pypi.org/simple" } }, ] sdist = { url = "https://files.pythonhosted.org/packages/7a/97/5a3609c4f8d58b039179648e62dd220f89864f56f7357f5d4f45c29eb2cc/scipy-1.17.1.tar.gz", hash = "sha256:95d8e012d8cb8816c226aef832200b1d45109ed4464303e997c5b13122b297c0", size = 30573822, upload-time = "2026-02-23T00:26:24.851Z" } wheels = [