From d9298b8bf51f8b1775cdd1b758ac0f2727e177b2 Mon Sep 17 00:00:00 2001 From: stephantul Date: Sat, 1 Aug 2026 11:10:42 +0200 Subject: [PATCH 1/9] feat: add integration tests --- Makefile | 15 +- model2vec/tokenizer/tokenizer.py | 8 +- tests/conftest.py | 3 +- tests/integration/__init__.py | 0 tests/integration/_distill_metrics.py | 104 +++++++++++ .../data/distilroberta_baseline.json | 128 ++++++++++++++ tests/integration/data/minilm_baseline.json | 128 ++++++++++++++ tests/integration/test_distill_regression.py | 162 ++++++++++++++++++ tests/integration/update_distill_baseline.py | 64 +++++++ 9 files changed, 608 insertions(+), 4 deletions(-) create mode 100644 tests/integration/__init__.py create mode 100644 tests/integration/_distill_metrics.py create mode 100644 tests/integration/data/distilroberta_baseline.json create mode 100644 tests/integration/data/minilm_baseline.json create mode 100644 tests/integration/test_distill_regression.py create mode 100644 tests/integration/update_distill_baseline.py diff --git a/Makefile b/Makefile index 276817ed..6c95d1c9 100644 --- a/Makefile +++ b/Makefile @@ -17,7 +17,20 @@ 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" + +# Downloads a real model, distills several variants from it, and compares the +# result (vocab size, token order, embedding rank, semantic sanity, etc) against +# the stored JSON baseline in tests/integration/data/distill_baseline.json. +# Not run as part of `make test`: it needs network access and is much slower. +test-integration: + uv run pytest tests/integration $(VERBOSITY) + +# Regenerates the JSON baseline used above. Only run this deliberately after a +# change that intentionally alters distillation output, then review the diff of +# tests/integration/data/distill_baseline.json before committing it. +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 c939a3bb..463f171b 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/tests/conftest.py b/tests/conftest.py index 2f45eef7..ba48e4f1 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,5 +1,6 @@ from __future__ import annotations +from types import SimpleNamespace from typing import Any, cast import numpy as np @@ -80,7 +81,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 = SimpleNamespace() 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 00000000..e69de29b diff --git a/tests/integration/_distill_metrics.py b/tests/integration/_distill_metrics.py new file mode 100644 index 00000000..1ef95152 --- /dev/null +++ b/tests/integration/_distill_metrics.py @@ -0,0 +1,104 @@ +"""Shared config and metric computation for the distillation regression test. + +Used by both `test_distill_regression.py` (which compares metrics against the stored baselines) +and `update_distill_baseline.py` (which (re)writes those baselines). +""" + +from __future__ import annotations + +import copy +import hashlib +from pathlib import Path +from typing import Any, cast + +import numpy as np +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: dict[str, str] = { + "minilm": "sentence-transformers/all-MiniLM-L6-v2", + "distilroberta": "sentence-transformers/all-distilroberta-v1", +} + +BASELINE_DIR = Path(__file__).parent / "data" + +_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"}, + "quantized": {"pca_dims": 256, "vocabulary_quantization": 2000, "quantize_to": "float32"}, +} + +SEMANTIC_TRIPLES = [ + ("king", "queen", "bicycle"), + ("dog", "puppy", "astronomy"), + ("happy", "joyful", "concrete"), +] + + +def baseline_path_for(model_name: str) -> Path: + """The JSON baseline file for a given short model name (a key of `BASE_MODELS`).""" + return BASELINE_DIR / f"{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=copy.deepcopy(model), tokenizer=tokenizer, **kwargs) + for name, kwargs in CONFIGS.items() + } + + +def _cosine_similarity(u: np.ndarray, v: np.ndarray) -> float: + return float(np.dot(u, v) / (np.linalg.norm(u) * np.linalg.norm(v))) + + +def compute_similarity_scores(model: StaticModel) -> dict[str, float]: + """Compute cosine similarity scores for the semantic sanity pairs.""" + scores: dict[str, float] = {} + for word_a, word_b, unrelated in SEMANTIC_TRIPLES: + vectors = model.encode([word_a, word_b, unrelated]) + scores[f"{word_a}~{word_b}"] = round(_cosine_similarity(vectors[0], vectors[1]), 6) + scores[f"{word_a}~{unrelated}"] = round(_cosine_similarity(vectors[0], vectors[2]), 6) + return scores + + +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 semantic + similarity 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:], + "similarity_scores": compute_similarity_scores(model), + } diff --git a/tests/integration/data/distilroberta_baseline.json b/tests/integration/data/distilroberta_baseline.json new file mode 100644 index 00000000..6721af5a --- /dev/null +++ b/tests/integration/data/distilroberta_baseline.json @@ -0,0 +1,128 @@ +{ + "base_model": "sentence-transformers/all-distilroberta-v1", + "configs": { + "custom_vocab": { + "embedding_dim": 32, + "embedding_mean": -0.0016, + "embedding_rank": 32, + "embedding_rows": 50277, + "embedding_std": 1.178123, + "first_tokens": [ + "", + "", + ".", + "\u0120the", + ",", + "\u0120to", + "\u0120and", + "\u0120of", + "\u0120a", + "\u0120in" + ], + "full_vocab_size": 50277, + "last_tokens": [ + "\u0120quixnorfle", + "norfle", + "\u0120quix", + "\u0120blorptastic", + "ptastic", + "\u0120blor", + "\u0120flimzycrag", + "zycr", + "\u0120flim", + "\u0120flimzycr" + ], + "similarity_scores": { + "dog~astronomy": -0.007912, + "dog~puppy": 0.74848, + "happy~concrete": -0.281206, + "happy~joyful": 0.400427, + "king~bicycle": -0.088169, + "king~queen": 0.749298 + }, + "token_order_hash": "5d377f1ee7affee5f45e4c3e1c222b5ef8204e80e88d2864e2fa24b134ca1e1e" + }, + "quantized": { + "embedding_dim": 256, + "embedding_mean": 0.0, + "embedding_rank": 256, + "embedding_rows": 2000, + "embedding_std": 0.516296, + "first_tokens": [ + "", + "", + ".", + "\u0120the", + ",", + "\u0120to", + "\u0120and", + "\u0120of", + "\u0120a", + "\u0120in" + ], + "full_vocab_size": 50262, + "last_tokens": [ + "", + "madeupword0000", + "madeupword0001", + "madeupword0002" + ], + "similarity_scores": { + "dog~astronomy": 0.004173, + "dog~puppy": 0.998192, + "happy~concrete": -0.087464, + "happy~joyful": 0.996333, + "king~bicycle": 0.023292, + "king~queen": 0.999984 + }, + "token_order_hash": "82f963ecc62f9a8c203ed70e8eadc2323389b82fc70b43a0ac614f5d76b8b479" + }, + "subword": { + "embedding_dim": 256, + "embedding_mean": -0.000478, + "embedding_rank": 256, + "embedding_rows": 50262, + "embedding_std": 0.649781, + "first_tokens": [ + "", + "", + ".", + "\u0120the", + ",", + "\u0120to", + "\u0120and", + "\u0120of", + "\u0120a", + "\u0120in" + ], + "full_vocab_size": 50262, + "last_tokens": [ + "", + "madeupword0000", + "madeupword0001", + "madeupword0002" + ], + "similarity_scores": { + "dog~astronomy": -0.064217, + "dog~puppy": 0.715823, + "happy~concrete": -0.102525, + "happy~joyful": 0.505929, + "king~bicycle": 0.01079, + "king~queen": 0.542192 + }, + "token_order_hash": "82f963ecc62f9a8c203ed70e8eadc2323389b82fc70b43a0ac614f5d76b8b479" + } + } +} diff --git a/tests/integration/data/minilm_baseline.json b/tests/integration/data/minilm_baseline.json new file mode 100644 index 00000000..e57cfdad --- /dev/null +++ b/tests/integration/data/minilm_baseline.json @@ -0,0 +1,128 @@ +{ + "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" + ], + "similarity_scores": { + "dog~astronomy": -0.301253, + "dog~puppy": 0.75149, + "happy~concrete": -0.033847, + "happy~joyful": 0.768889, + "king~bicycle": 0.19222, + "king~queen": 0.568608 + }, + "token_order_hash": "09d9aaaedc7d41b2fd1e3ffa64f16794a2cefe9ffc7faf1055ae5d49350c3060" + }, + "quantized": { + "embedding_dim": 256, + "embedding_mean": -0.0, + "embedding_rank": 256, + "embedding_rows": 2000, + "embedding_std": 0.304732, + "first_tokens": [ + "[PAD]", + "[UNK]", + "!", + "\"", + "#", + "$", + "%", + "&", + "'", + "(" + ], + "full_vocab_size": 29525, + "last_tokens": [ + "##\uff01", + "##\uff08", + "##\uff09", + "##\uff0c", + "##\uff0d", + "##\uff0e", + "##\uff0f", + "##\uff1a", + "##\uff1f", + "##\uff5e" + ], + "similarity_scores": { + "dog~astronomy": -0.079069, + "dog~puppy": 0.998789, + "happy~concrete": -0.108555, + "happy~joyful": 0.828377, + "king~bicycle": 0.043042, + "king~queen": 0.229007 + }, + "token_order_hash": "d881fab650dd73240615f87f531ff31a6135f961933ebc32c0dfdd130763e7a8" + }, + "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" + ], + "similarity_scores": { + "dog~astronomy": -0.04307, + "dog~puppy": 0.716282, + "happy~concrete": -0.002784, + "happy~joyful": 0.585774, + "king~bicycle": 0.075807, + "king~queen": 0.57314 + }, + "token_order_hash": "d881fab650dd73240615f87f531ff31a6135f961933ebc32c0dfdd130763e7a8" + } + } +} diff --git a/tests/integration/test_distill_regression.py b/tests/integration/test_distill_regression.py new file mode 100644 index 00000000..eca5fc6e --- /dev/null +++ b/tests/integration/test_distill_regression.py @@ -0,0 +1,162 @@ +"""Regression test for the distillation pipeline. + +Downloads real sentence-transformers (see `BASE_MODELS`), distills several StaticModel variants +from each, and compares the result against a stored JSON baseline +(`tests/integration/data/_baseline.json`): vocabulary size, token order, embedding matrix +rank/distribution, and semantic similarity scores. This catches regressions that unit tests (which +run against a mocked transformer) can't, e.g. a tokenizer/embedding misalignment, a degenerate +(rank-collapsed) embedding matrix, or a silent drop in semantic quality -- and, by covering more +than one base model, regressions specific to a particular tokenizer family (e.g. BPE vs. wordpiece). + +This suite requires network access and real forward passes over a full vocabulary, so it is +intentionally excluded from `make test` and must be run explicitly with `make test-integration`. + +If a change intentionally alters distillation output, regenerate the baselines with +`make test-integration-update` and review the JSON diff before committing it. +""" + +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(BASE_MODELS[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_semantic_similarity_matches_baseline( + config_name: str, model_name: str, baseline: dict[str, Any], current_metrics: dict[str, dict[str, Any]] +) -> None: + """Guard against silently destroyed performance.""" + expected_scores = baseline["configs"][config_name]["similarity_scores"] + actual_scores = current_metrics[config_name]["similarity_scores"] + for pair_name, expected_score in expected_scores.items(): + actual_score = actual_scores[pair_name] + assert actual_score == pytest.approx(expected_score, abs=0.05), ( + f"[{model_name}/{config_name}] similarity for '{pair_name}' 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 00000000..ecf9b14e --- /dev/null +++ b/tests/integration/update_distill_baseline.py @@ -0,0 +1,64 @@ +"""Regenerate the golden baselines used by `test_distill_regression.py`. + +Run this deliberately (`make test-integration-update`) after a change that intentionally alters +distillation output (e.g. a new default, a bugfix that shifts embeddings). Then inspect the JSON +diff (`git diff tests/integration/data/`) to confirm the change is expected before committing it -- +an unreviewed update here would hide a real regression. + +By default this regenerates the baseline for every model in `BASE_MODELS`. Pass one or more short +names (the keys of `BASE_MODELS`, e.g. `minilm`) to only regenerate those, e.g.: + + uv run python -m tests.integration.update_distill_baseline minilm +""" + +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 short name, i.e. a key of `BASE_MODELS`. + """ + hub_model_name = BASE_MODELS[model_name] + model, tokenizer = load_base_model_and_tokenizer(hub_model_name) + distilled = distill_all(model, tokenizer) + + baseline = { + "base_model": hub_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}' ({hub_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) + unknown = [name for name in requested if name not in BASE_MODELS] + if unknown: + raise SystemExit(f"Unknown model name(s) {unknown}. Choose from {sorted(BASE_MODELS)}.") + + for model_name in requested: + update_baseline(model_name) + + +if __name__ == "__main__": + logging.basicConfig(level=logging.INFO, format="%(message)s") + main() From b23cb3b923673cb8e48abe26ed486af99563f8a6 Mon Sep 17 00:00:00 2001 From: stephantul Date: Sat, 1 Aug 2026 11:43:38 +0200 Subject: [PATCH 2/9] temporarily relax skeletoken to 0.3.3 --- pyproject.toml | 2 +- uv.lock | 90 +++++++++++++++++++++++++------------------------- 2 files changed, 46 insertions(+), 46 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index dd8cf83d..54e5ad5d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -58,7 +58,7 @@ dev = [ "setuptools", ] -distill = ["torch", "transformers<5.4.0", "scikit-learn", "skeletoken>=0.3.3"] +distill = ["torch", "transformers<5.4.0", "scikit-learn", "skeletoken>=0.3.0,<0.5.0"] onnx = ["onnx", "torch"] # train also installs inference train = ["torch", "lightning", "scikit-learn", "skops"] diff --git a/uv.lock b/uv.lock index 000d5363..e76412cb 100644 --- a/uv.lock +++ b/uv.lock @@ -418,7 +418,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 +702,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 +727,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 +753,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 +774,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 +1168,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.3.0,<0.5.0" }, { name = "skops", marker = "extra == 'inference'" }, { name = "skops", marker = "extra == 'train'" }, { name = "tokenizers", specifier = ">=0.20" }, @@ -2444,10 +2444,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 +2493,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 +2546,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 +2607,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 = [ From 456545648f85318844eb316e5a35d0e92f6b533e Mon Sep 17 00:00:00 2001 From: stephantul Date: Sat, 1 Aug 2026 13:39:01 +0200 Subject: [PATCH 3/9] add newer models --- Makefile | 7 - tests/integration/_distill_metrics.py | 27 ++-- ...ba-NLP___gte-modernbert-base_baseline.json | 87 ++++++++++++ .../baai___bge-base-en-v1.5_baseline.json | 87 ++++++++++++ .../data/distilroberta_baseline.json | 128 ------------------ ...google___embeddinggemma-300m_baseline.json | 87 ++++++++++++ ...float___multilingual-e5-base_baseline.json | 87 ++++++++++++ ...sformers___all-MiniLM-L6-v2_baseline.json} | 41 ------ tests/integration/update_distill_baseline.py | 23 +--- 9 files changed, 361 insertions(+), 213 deletions(-) create mode 100644 tests/integration/data/Alibaba-NLP___gte-modernbert-base_baseline.json create mode 100644 tests/integration/data/baai___bge-base-en-v1.5_baseline.json delete mode 100644 tests/integration/data/distilroberta_baseline.json create mode 100644 tests/integration/data/google___embeddinggemma-300m_baseline.json create mode 100644 tests/integration/data/intfloat___multilingual-e5-base_baseline.json rename tests/integration/data/{minilm_baseline.json => sentence-transformers___all-MiniLM-L6-v2_baseline.json} (67%) diff --git a/Makefile b/Makefile index 6c95d1c9..083b4474 100644 --- a/Makefile +++ b/Makefile @@ -22,15 +22,8 @@ test: test-verbose: make test VERBOSITY="-vvv" -# Downloads a real model, distills several variants from it, and compares the -# result (vocab size, token order, embedding rank, semantic sanity, etc) against -# the stored JSON baseline in tests/integration/data/distill_baseline.json. -# Not run as part of `make test`: it needs network access and is much slower. test-integration: uv run pytest tests/integration $(VERBOSITY) -# Regenerates the JSON baseline used above. Only run this deliberately after a -# change that intentionally alters distillation output, then review the diff of -# tests/integration/data/distill_baseline.json before committing it. test-integration-update: uv run python -m tests.integration.update_distill_baseline diff --git a/tests/integration/_distill_metrics.py b/tests/integration/_distill_metrics.py index 1ef95152..607a4d51 100644 --- a/tests/integration/_distill_metrics.py +++ b/tests/integration/_distill_metrics.py @@ -1,12 +1,5 @@ -"""Shared config and metric computation for the distillation regression test. - -Used by both `test_distill_regression.py` (which compares metrics against the stored baselines) -and `update_distill_baseline.py` (which (re)writes those baselines). -""" - from __future__ import annotations -import copy import hashlib from pathlib import Path from typing import Any, cast @@ -19,10 +12,13 @@ from model2vec.distill import distill_from_model from model2vec.model import StaticModel -BASE_MODELS: dict[str, str] = { - "minilm": "sentence-transformers/all-MiniLM-L6-v2", - "distilroberta": "sentence-transformers/all-distilroberta-v1", -} +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" @@ -31,7 +27,6 @@ CONFIGS: dict[str, dict[str, Any]] = { "subword": {"pca_dims": 256, "quantize_to": "float32"}, "custom_vocab": {"vocabulary": _NOVEL_VOCABULARY, "pca_dims": 32, "quantize_to": "float32"}, - "quantized": {"pca_dims": 256, "vocabulary_quantization": 2000, "quantize_to": "float32"}, } SEMANTIC_TRIPLES = [ @@ -43,7 +38,8 @@ def baseline_path_for(model_name: str) -> Path: """The JSON baseline file for a given short model name (a key of `BASE_MODELS`).""" - return BASELINE_DIR / f"{model_name}_baseline.json" + 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]: @@ -59,10 +55,7 @@ def load_base_model_and_tokenizer(model_name: str) -> tuple[PreTrainedModel, Pre 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=copy.deepcopy(model), tokenizer=tokenizer, **kwargs) - for name, kwargs in CONFIGS.items() - } + return {name: distill_from_model(model=model, tokenizer=tokenizer, **kwargs) for name, kwargs in CONFIGS.items()} def _cosine_similarity(u: np.ndarray, v: np.ndarray) -> float: 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 00000000..adf9fa44 --- /dev/null +++ b/tests/integration/data/Alibaba-NLP___gte-modernbert-base_baseline.json @@ -0,0 +1,87 @@ +{ + "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" + ], + "similarity_scores": { + "dog~astronomy": -0.336587, + "dog~puppy": -0.012992, + "happy~concrete": 0.197486, + "happy~joyful": 0.085868, + "king~bicycle": 0.135328, + "king~queen": 0.257917 + }, + "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]" + ], + "similarity_scores": { + "dog~astronomy": -0.19247, + "dog~puppy": 0.069406, + "happy~concrete": 0.101801, + "happy~joyful": 0.106968, + "king~bicycle": 0.030125, + "king~queen": 0.21826 + }, + "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 00000000..abf4633c --- /dev/null +++ b/tests/integration/data/baai___bge-base-en-v1.5_baseline.json @@ -0,0 +1,87 @@ +{ + "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" + ], + "similarity_scores": { + "dog~astronomy": -0.059722, + "dog~puppy": 0.507191, + "happy~concrete": -0.047596, + "happy~joyful": 0.78852, + "king~bicycle": 0.226473, + "king~queen": 0.484422 + }, + "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" + ], + "similarity_scores": { + "dog~astronomy": 0.043488, + "dog~puppy": 0.718679, + "happy~concrete": -0.00541, + "happy~joyful": 0.603683, + "king~bicycle": 0.103374, + "king~queen": 0.406895 + }, + "token_order_hash": "d881fab650dd73240615f87f531ff31a6135f961933ebc32c0dfdd130763e7a8" + } + } +} diff --git a/tests/integration/data/distilroberta_baseline.json b/tests/integration/data/distilroberta_baseline.json deleted file mode 100644 index 6721af5a..00000000 --- a/tests/integration/data/distilroberta_baseline.json +++ /dev/null @@ -1,128 +0,0 @@ -{ - "base_model": "sentence-transformers/all-distilroberta-v1", - "configs": { - "custom_vocab": { - "embedding_dim": 32, - "embedding_mean": -0.0016, - "embedding_rank": 32, - "embedding_rows": 50277, - "embedding_std": 1.178123, - "first_tokens": [ - "", - "", - ".", - "\u0120the", - ",", - "\u0120to", - "\u0120and", - "\u0120of", - "\u0120a", - "\u0120in" - ], - "full_vocab_size": 50277, - "last_tokens": [ - "\u0120quixnorfle", - "norfle", - "\u0120quix", - "\u0120blorptastic", - "ptastic", - "\u0120blor", - "\u0120flimzycrag", - "zycr", - "\u0120flim", - "\u0120flimzycr" - ], - "similarity_scores": { - "dog~astronomy": -0.007912, - "dog~puppy": 0.74848, - "happy~concrete": -0.281206, - "happy~joyful": 0.400427, - "king~bicycle": -0.088169, - "king~queen": 0.749298 - }, - "token_order_hash": "5d377f1ee7affee5f45e4c3e1c222b5ef8204e80e88d2864e2fa24b134ca1e1e" - }, - "quantized": { - "embedding_dim": 256, - "embedding_mean": 0.0, - "embedding_rank": 256, - "embedding_rows": 2000, - "embedding_std": 0.516296, - "first_tokens": [ - "", - "", - ".", - "\u0120the", - ",", - "\u0120to", - "\u0120and", - "\u0120of", - "\u0120a", - "\u0120in" - ], - "full_vocab_size": 50262, - "last_tokens": [ - "", - "madeupword0000", - "madeupword0001", - "madeupword0002" - ], - "similarity_scores": { - "dog~astronomy": 0.004173, - "dog~puppy": 0.998192, - "happy~concrete": -0.087464, - "happy~joyful": 0.996333, - "king~bicycle": 0.023292, - "king~queen": 0.999984 - }, - "token_order_hash": "82f963ecc62f9a8c203ed70e8eadc2323389b82fc70b43a0ac614f5d76b8b479" - }, - "subword": { - "embedding_dim": 256, - "embedding_mean": -0.000478, - "embedding_rank": 256, - "embedding_rows": 50262, - "embedding_std": 0.649781, - "first_tokens": [ - "", - "", - ".", - "\u0120the", - ",", - "\u0120to", - "\u0120and", - "\u0120of", - "\u0120a", - "\u0120in" - ], - "full_vocab_size": 50262, - "last_tokens": [ - "", - "madeupword0000", - "madeupword0001", - "madeupword0002" - ], - "similarity_scores": { - "dog~astronomy": -0.064217, - "dog~puppy": 0.715823, - "happy~concrete": -0.102525, - "happy~joyful": 0.505929, - "king~bicycle": 0.01079, - "king~queen": 0.542192 - }, - "token_order_hash": "82f963ecc62f9a8c203ed70e8eadc2323389b82fc70b43a0ac614f5d76b8b479" - } - } -} 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 00000000..7640539c --- /dev/null +++ b/tests/integration/data/google___embeddinggemma-300m_baseline.json @@ -0,0 +1,87 @@ +{ + "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" + ], + "similarity_scores": { + "dog~astronomy": -0.004148, + "dog~puppy": 0.421932, + "happy~concrete": -0.209418, + "happy~joyful": 0.94543, + "king~bicycle": -0.03067, + "king~queen": 0.68369 + }, + "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" + ], + "similarity_scores": { + "dog~astronomy": 0.032034, + "dog~puppy": 0.218597, + "happy~concrete": -0.091707, + "happy~joyful": 0.88989, + "king~bicycle": -0.061335, + "king~queen": 0.492228 + }, + "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 00000000..14ab03a1 --- /dev/null +++ b/tests/integration/data/intfloat___multilingual-e5-base_baseline.json @@ -0,0 +1,87 @@ +{ + "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" + ], + "similarity_scores": { + "dog~astronomy": 0.573303, + "dog~puppy": 0.880382, + "happy~concrete": 0.520932, + "happy~joyful": 0.807798, + "king~bicycle": 0.660403, + "king~queen": 0.535522 + }, + "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" + ], + "similarity_scores": { + "dog~astronomy": 0.257133, + "dog~puppy": 0.777345, + "happy~concrete": 0.21136, + "happy~joyful": 0.663768, + "king~bicycle": 0.367408, + "king~queen": 0.285096 + }, + "token_order_hash": "d0534dadde49bec613f7090e0fec5797f10194566612acf36895214252e290e0" + } + } +} diff --git a/tests/integration/data/minilm_baseline.json b/tests/integration/data/sentence-transformers___all-MiniLM-L6-v2_baseline.json similarity index 67% rename from tests/integration/data/minilm_baseline.json rename to tests/integration/data/sentence-transformers___all-MiniLM-L6-v2_baseline.json index e57cfdad..3d8e73aa 100644 --- a/tests/integration/data/minilm_baseline.json +++ b/tests/integration/data/sentence-transformers___all-MiniLM-L6-v2_baseline.json @@ -42,47 +42,6 @@ }, "token_order_hash": "09d9aaaedc7d41b2fd1e3ffa64f16794a2cefe9ffc7faf1055ae5d49350c3060" }, - "quantized": { - "embedding_dim": 256, - "embedding_mean": -0.0, - "embedding_rank": 256, - "embedding_rows": 2000, - "embedding_std": 0.304732, - "first_tokens": [ - "[PAD]", - "[UNK]", - "!", - "\"", - "#", - "$", - "%", - "&", - "'", - "(" - ], - "full_vocab_size": 29525, - "last_tokens": [ - "##\uff01", - "##\uff08", - "##\uff09", - "##\uff0c", - "##\uff0d", - "##\uff0e", - "##\uff0f", - "##\uff1a", - "##\uff1f", - "##\uff5e" - ], - "similarity_scores": { - "dog~astronomy": -0.079069, - "dog~puppy": 0.998789, - "happy~concrete": -0.108555, - "happy~joyful": 0.828377, - "king~bicycle": 0.043042, - "king~queen": 0.229007 - }, - "token_order_hash": "d881fab650dd73240615f87f531ff31a6135f961933ebc32c0dfdd130763e7a8" - }, "subword": { "embedding_dim": 256, "embedding_mean": -0.0001, diff --git a/tests/integration/update_distill_baseline.py b/tests/integration/update_distill_baseline.py index ecf9b14e..b2bc4725 100644 --- a/tests/integration/update_distill_baseline.py +++ b/tests/integration/update_distill_baseline.py @@ -1,16 +1,3 @@ -"""Regenerate the golden baselines used by `test_distill_regression.py`. - -Run this deliberately (`make test-integration-update`) after a change that intentionally alters -distillation output (e.g. a new default, a bugfix that shifts embeddings). Then inspect the JSON -diff (`git diff tests/integration/data/`) to confirm the change is expected before committing it -- -an unreviewed update here would hide a real regression. - -By default this regenerates the baseline for every model in `BASE_MODELS`. Pass one or more short -names (the keys of `BASE_MODELS`, e.g. `minilm`) to only regenerate those, e.g.: - - uv run python -m tests.integration.update_distill_baseline minilm -""" - from __future__ import annotations import json @@ -33,27 +20,23 @@ def update_baseline(model_name: str) -> None: :param model_name: A short name, i.e. a key of `BASE_MODELS`. """ - hub_model_name = BASE_MODELS[model_name] - model, tokenizer = load_base_model_and_tokenizer(hub_model_name) + model, tokenizer = load_base_model_and_tokenizer(model_name) distilled = distill_all(model, tokenizer) baseline = { - "base_model": hub_model_name, + "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}' ({hub_model_name}) to {path}") + 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) - unknown = [name for name in requested if name not in BASE_MODELS] - if unknown: - raise SystemExit(f"Unknown model name(s) {unknown}. Choose from {sorted(BASE_MODELS)}.") for model_name in requested: update_baseline(model_name) From 206f183b77af3f492f0f43be6ca4fcd61f85f139 Mon Sep 17 00:00:00 2001 From: stephantul Date: Sat, 1 Aug 2026 13:47:03 +0200 Subject: [PATCH 4/9] fix test path --- tests/integration/test_distill_regression.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/integration/test_distill_regression.py b/tests/integration/test_distill_regression.py index eca5fc6e..3ee38614 100644 --- a/tests/integration/test_distill_regression.py +++ b/tests/integration/test_distill_regression.py @@ -66,7 +66,7 @@ def baseline(model_name: str) -> dict[str, Any]: @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(BASE_MODELS[model_name]) + return load_base_model_and_tokenizer(model_name) @pytest.fixture(scope="module") From b5bae433db1ae28d6be4ab18ab030b5fc0c9a0f8 Mon Sep 17 00:00:00 2001 From: stephantul Date: Sat, 1 Aug 2026 13:48:15 +0200 Subject: [PATCH 5/9] ignore integration tests in CI --- .github/workflows/ci.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index 5ded6cd3..cf9224fb 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 From 326aa031c8f7aef6e453f3945be6c76776d6c9cb Mon Sep 17 00:00:00 2001 From: stephantul Date: Sat, 1 Aug 2026 13:59:16 +0200 Subject: [PATCH 6/9] fix test --- tests/conftest.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/tests/conftest.py b/tests/conftest.py index ba48e4f1..af463433 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,7 +1,6 @@ from __future__ import annotations -from types import SimpleNamespace -from typing import Any, cast +from typing import Any, Iterator, cast import numpy as np import pytest @@ -63,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.""" @@ -81,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 = SimpleNamespace() + self.config = DumbConfig() def to(self, device: str) -> MockPreTrainedModel: self.device = device From d653c3f95b5c0978179282342f416916abafa997 Mon Sep 17 00:00:00 2001 From: stephantul Date: Sat, 1 Aug 2026 14:02:50 +0200 Subject: [PATCH 7/9] fix skeletoken --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 54e5ad5d..19799af6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -58,7 +58,7 @@ dev = [ "setuptools", ] -distill = ["torch", "transformers<5.4.0", "scikit-learn", "skeletoken>=0.3.0,<0.5.0"] +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"] From 055ff3b9f5efcb9b480023af5eb488602a475f37 Mon Sep 17 00:00:00 2001 From: stephantul Date: Sat, 1 Aug 2026 14:15:07 +0200 Subject: [PATCH 8/9] update specifier --- pyproject.toml | 1 + uv.lock | 5 ++++- 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 19799af6..f9dcbd95 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -123,3 +123,4 @@ version = {attr = "model2vec.version.__version__"} [tool.uv] exclude-newer = "1 week" +exclude-newer-package = { skeletoken = false } diff --git a/uv.lock b/uv.lock index e76412cb..28f802b6 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" @@ -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.0,<0.5.0" }, + { 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" }, From d9a36b57c39756acef8960ea7e86de054daf3b93 Mon Sep 17 00:00:00 2001 From: stephantul Date: Sat, 1 Aug 2026 18:30:30 +0200 Subject: [PATCH 9/9] add MTEB as dependency --- Makefile | 2 +- pyproject.toml | 1 + ...ba-NLP___gte-modernbert-base_baseline.json | 32 ++++++---- .../baai___bge-base-en-v1.5_baseline.json | 32 ++++++---- ...google___embeddinggemma-300m_baseline.json | 32 ++++++---- ...float___multilingual-e5-base_baseline.json | 32 ++++++---- ...nsformers___all-MiniLM-L6-v2_baseline.json | 32 ++++++---- ..._distill_metrics.py => distill_metrics.py} | 64 +++++++++++++------ tests/integration/test_distill_regression.py | 35 +++------- tests/integration/update_distill_baseline.py | 4 +- 10 files changed, 148 insertions(+), 118 deletions(-) rename tests/integration/{_distill_metrics.py => distill_metrics.py} (65%) diff --git a/Makefile b/Makefile index 083b4474..c067b773 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 diff --git a/pyproject.toml b/pyproject.toml index f9dcbd95..0a9e75e2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -64,6 +64,7 @@ onnx = ["onnx", "torch"] train = ["torch", "lightning", "scikit-learn", "skops"] inference = ["scikit-learn", "skops"] quantization = ["scikit-learn"] +integration = ["mteb"] [project.urls] "Homepage" = "https://github.com/MinishLab" diff --git a/tests/integration/data/Alibaba-NLP___gte-modernbert-base_baseline.json b/tests/integration/data/Alibaba-NLP___gte-modernbert-base_baseline.json index adf9fa44..0f10306d 100644 --- a/tests/integration/data/Alibaba-NLP___gte-modernbert-base_baseline.json +++ b/tests/integration/data/Alibaba-NLP___gte-modernbert-base_baseline.json @@ -32,13 +32,15 @@ "\u0120flim", "\u0120flimzycr" ], - "similarity_scores": { - "dog~astronomy": -0.336587, - "dog~puppy": -0.012992, - "happy~concrete": 0.197486, - "happy~joyful": 0.085868, - "king~bicycle": 0.135328, - "king~queen": 0.257917 + "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" }, @@ -73,13 +75,15 @@ "[UNK]", "[PAD]" ], - "similarity_scores": { - "dog~astronomy": -0.19247, - "dog~puppy": 0.069406, - "happy~concrete": 0.101801, - "happy~joyful": 0.106968, - "king~bicycle": 0.030125, - "king~queen": 0.21826 + "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 index abf4633c..da97ebf8 100644 --- a/tests/integration/data/baai___bge-base-en-v1.5_baseline.json +++ b/tests/integration/data/baai___bge-base-en-v1.5_baseline.json @@ -32,13 +32,15 @@ "blorptastic", "flimzycrag" ], - "similarity_scores": { - "dog~astronomy": -0.059722, - "dog~puppy": 0.507191, - "happy~concrete": -0.047596, - "happy~joyful": 0.78852, - "king~bicycle": 0.226473, - "king~queen": 0.484422 + "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" }, @@ -73,13 +75,15 @@ "##\uff1f", "##\uff5e" ], - "similarity_scores": { - "dog~astronomy": 0.043488, - "dog~puppy": 0.718679, - "happy~concrete": -0.00541, - "happy~joyful": 0.603683, - "king~bicycle": 0.103374, - "king~queen": 0.406895 + "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 index 7640539c..e1003ece 100644 --- a/tests/integration/data/google___embeddinggemma-300m_baseline.json +++ b/tests/integration/data/google___embeddinggemma-300m_baseline.json @@ -32,13 +32,15 @@ "flimzycr", "zycr" ], - "similarity_scores": { - "dog~astronomy": -0.004148, - "dog~puppy": 0.421932, - "happy~concrete": -0.209418, - "happy~joyful": 0.94543, - "king~bicycle": -0.03067, - "king~queen": 0.68369 + "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" }, @@ -73,13 +75,15 @@ "\ud44e", "\ud835\udd39" ], - "similarity_scores": { - "dog~astronomy": 0.032034, - "dog~puppy": 0.218597, - "happy~concrete": -0.091707, - "happy~joyful": 0.88989, - "king~bicycle": -0.061335, - "king~queen": 0.492228 + "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 index 14ab03a1..d2299a5e 100644 --- a/tests/integration/data/intfloat___multilingual-e5-base_baseline.json +++ b/tests/integration/data/intfloat___multilingual-e5-base_baseline.json @@ -32,13 +32,15 @@ "\u2581blorptastic", "\u2581flimzycrag" ], - "similarity_scores": { - "dog~astronomy": 0.573303, - "dog~puppy": 0.880382, - "happy~concrete": 0.520932, - "happy~joyful": 0.807798, - "king~bicycle": 0.660403, - "king~queen": 0.535522 + "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" }, @@ -73,13 +75,15 @@ "\u1ea5", "\u7a23" ], - "similarity_scores": { - "dog~astronomy": 0.257133, - "dog~puppy": 0.777345, - "happy~concrete": 0.21136, - "happy~joyful": 0.663768, - "king~bicycle": 0.367408, - "king~queen": 0.285096 + "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 index 3d8e73aa..50a54a0e 100644 --- a/tests/integration/data/sentence-transformers___all-MiniLM-L6-v2_baseline.json +++ b/tests/integration/data/sentence-transformers___all-MiniLM-L6-v2_baseline.json @@ -32,13 +32,15 @@ "blorptastic", "flimzycrag" ], - "similarity_scores": { - "dog~astronomy": -0.301253, - "dog~puppy": 0.75149, - "happy~concrete": -0.033847, - "happy~joyful": 0.768889, - "king~bicycle": 0.19222, - "king~queen": 0.568608 + "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" }, @@ -73,13 +75,15 @@ "##\uff1f", "##\uff5e" ], - "similarity_scores": { - "dog~astronomy": -0.04307, - "dog~puppy": 0.716282, - "happy~concrete": -0.002784, - "happy~joyful": 0.585774, - "king~bicycle": 0.075807, - "king~queen": 0.57314 + "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 similarity index 65% rename from tests/integration/_distill_metrics.py rename to tests/integration/distill_metrics.py index 607a4d51..7b8ee4c0 100644 --- a/tests/integration/_distill_metrics.py +++ b/tests/integration/distill_metrics.py @@ -4,7 +4,10 @@ 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 @@ -22,6 +25,17 @@ 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]] = { @@ -29,12 +43,6 @@ "custom_vocab": {"vocabulary": _NOVEL_VOCABULARY, "pca_dims": 32, "quantize_to": "float32"}, } -SEMANTIC_TRIPLES = [ - ("king", "queen", "bicycle"), - ("dog", "puppy", "astronomy"), - ("happy", "joyful", "concrete"), -] - def baseline_path_for(model_name: str) -> Path: """The JSON baseline file for a given short model name (a key of `BASE_MODELS`).""" @@ -58,26 +66,44 @@ def distill_all(model: PreTrainedModel, tokenizer: PreTrainedTokenizerFast) -> d return {name: distill_from_model(model=model, tokenizer=tokenizer, **kwargs) for name, kwargs in CONFIGS.items()} -def _cosine_similarity(u: np.ndarray, v: np.ndarray) -> float: - return float(np.dot(u, v) / (np.linalg.norm(u) * np.linalg.norm(v))) +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_similarity_scores(model: StaticModel) -> dict[str, float]: - """Compute cosine similarity scores for the semantic sanity pairs.""" - scores: dict[str, float] = {} - for word_a, word_b, unrelated in SEMANTIC_TRIPLES: - vectors = model.encode([word_a, word_b, unrelated]) - scores[f"{word_a}~{word_b}"] = round(_cosine_similarity(vectors[0], vectors[1]), 6) - scores[f"{word_a}~{unrelated}"] = round(_cosine_similarity(vectors[0], vectors[2]), 6) - return scores + +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 semantic - similarity scores. Used both to write and to check the regression baseline. + :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) @@ -93,5 +119,5 @@ def compute_metrics(model: StaticModel) -> dict[str, Any]: "token_order_hash": token_order_hash, "first_tokens": tokens[:10], "last_tokens": tokens[-10:], - "similarity_scores": compute_similarity_scores(model), + "mteb_sts_scores": compute_mteb_sts_scores(model), } diff --git a/tests/integration/test_distill_regression.py b/tests/integration/test_distill_regression.py index 3ee38614..6d178d50 100644 --- a/tests/integration/test_distill_regression.py +++ b/tests/integration/test_distill_regression.py @@ -1,20 +1,3 @@ -"""Regression test for the distillation pipeline. - -Downloads real sentence-transformers (see `BASE_MODELS`), distills several StaticModel variants -from each, and compares the result against a stored JSON baseline -(`tests/integration/data/_baseline.json`): vocabulary size, token order, embedding matrix -rank/distribution, and semantic similarity scores. This catches regressions that unit tests (which -run against a mocked transformer) can't, e.g. a tokenizer/embedding misalignment, a degenerate -(rank-collapsed) embedding matrix, or a silent drop in semantic quality -- and, by covering more -than one base model, regressions specific to a particular tokenizer family (e.g. BPE vs. wordpiece). - -This suite requires network access and real forward passes over a full vocabulary, so it is -intentionally excluded from `make test` and must be run explicitly with `make test-integration`. - -If a change intentionally alters distillation output, regenerate the baselines with -`make test-integration-update` and review the JSON diff before committing it. -""" - from __future__ import annotations import copy @@ -29,7 +12,7 @@ from model2vec.distill import distill_from_model from model2vec.model import StaticModel -from tests.integration._distill_metrics import ( +from tests.integration.distill_metrics import ( BASE_MODELS, CONFIGS, baseline_path_for, @@ -123,16 +106,16 @@ def test_embedding_distribution_matches_baseline( @pytest.mark.parametrize("config_name", sorted(CONFIGS)) -def test_semantic_similarity_matches_baseline( +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: - """Guard against silently destroyed performance.""" - expected_scores = baseline["configs"][config_name]["similarity_scores"] - actual_scores = current_metrics[config_name]["similarity_scores"] - for pair_name, expected_score in expected_scores.items(): - actual_score = actual_scores[pair_name] - assert actual_score == pytest.approx(expected_score, abs=0.05), ( - f"[{model_name}/{config_name}] similarity for '{pair_name}' drifted from baseline: " + """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}" ) diff --git a/tests/integration/update_distill_baseline.py b/tests/integration/update_distill_baseline.py index b2bc4725..8a0dee9f 100644 --- a/tests/integration/update_distill_baseline.py +++ b/tests/integration/update_distill_baseline.py @@ -4,7 +4,7 @@ import logging import sys -from tests.integration._distill_metrics import ( +from tests.integration.distill_metrics import ( BASE_MODELS, baseline_path_for, compute_metrics, @@ -18,7 +18,7 @@ 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 short name, i.e. a key of `BASE_MODELS`. + :param model_name: A model identifier on Hugging Face. """ model, tokenizer = load_base_model_and_tokenizer(model_name) distilled = distill_all(model, tokenizer)