From 6cb535b9ea0b2773c51789c247dc4ca61fed8ea1 Mon Sep 17 00:00:00 2001 From: stephantul Date: Fri, 28 Aug 2026 17:51:06 +0200 Subject: [PATCH 1/2] feat: new integration test --- Makefile | 3 + .../minishlab___potion-base-32m_baseline.json | 73 ++++++++ .../minishlab___potion-base-8m_baseline.json | 73 ++++++++ ...b___potion-multilingual-128m_baseline.json | 177 ++++++++++++++++++ ...shlab___potion-retrieval-32m_baseline.json | 73 ++++++++ tests/integration/pretrained_model_metrics.py | 113 +++++++++++ .../test_pretrained_model_regression.py | 160 ++++++++++++++++ .../integration/update_pretrained_baseline.py | 41 ++++ 8 files changed, 713 insertions(+) create mode 100644 tests/integration/data/pretrained/minishlab___potion-base-32m_baseline.json create mode 100644 tests/integration/data/pretrained/minishlab___potion-base-8m_baseline.json create mode 100644 tests/integration/data/pretrained/minishlab___potion-multilingual-128m_baseline.json create mode 100644 tests/integration/data/pretrained/minishlab___potion-retrieval-32m_baseline.json create mode 100644 tests/integration/pretrained_model_metrics.py create mode 100644 tests/integration/test_pretrained_model_regression.py create mode 100644 tests/integration/update_pretrained_baseline.py diff --git a/Makefile b/Makefile index cc1db1a..0b03d1d 100644 --- a/Makefile +++ b/Makefile @@ -27,3 +27,6 @@ test-integration: test-integration-update: uv run python -m tests.integration.update_distill_baseline + +test-integration-pretrained-update: + uv run python -m tests.integration.update_pretrained_baseline diff --git a/tests/integration/data/pretrained/minishlab___potion-base-32m_baseline.json b/tests/integration/data/pretrained/minishlab___potion-base-32m_baseline.json new file mode 100644 index 0000000..8a29850 --- /dev/null +++ b/tests/integration/data/pretrained/minishlab___potion-base-32m_baseline.json @@ -0,0 +1,73 @@ +{ + "metrics": { + "base_model_name": null, + "config": { + "apply_pca": 512, + "apply_zipf": true, + "architectures": [ + "StaticModel" + ], + "hidden_dim": 512, + "model_type": "model2vec", + "normalize": true, + "seq_length": 1000000, + "tokenizer_name": "baai/bge-base-en-v1.5" + }, + "embedding_dim": 512, + "embedding_dtype": "float32", + "embedding_mean": -0.000142, + "embedding_rank": 512, + "embedding_row_norm_mean": 109.413377, + "embedding_rows": 63091, + "embedding_std": 4.876269, + "encoding_speed": { + "sentences_per_second": 119227.47, + "tokens_per_second": 1341309.0 + }, + "first_tokens": [ + "[PAD]", + "[UNK]", + "[CLS]", + "[SEP]", + "[MASK]", + "!", + "\"", + "#", + "$", + "%" + ], + "full_vocab_size": 63091, + "has_token_mapping": false, + "has_weights": false, + "language": null, + "last_tokens": [ + "nti", + "corporeal", + "downgrading", + "symbolise", + "emptive", + "daters", + "cousteau", + "crichton", + "60k", + "longines" + ], + "median_token_length": 7, + "mteb_sts_scores": { + "BIOSSES": 0.775566, + "SICK-R": 0.656684, + "STS12": 0.627234, + "STS13": 0.775936, + "STS14": 0.729035, + "STS15": 0.807592, + "STS16": 0.768566, + "STSBenchmark": 0.768385 + }, + "normalize": true, + "token_order_hash": "ba90854d713d66960327011c1d592b5f16c6dd7fe26676b103490b62753dc7c9", + "tokenizer_type": "WordPiece", + "unk_token_id": 1, + "vocabulary_quantization": null + }, + "model": "minishlab/potion-base-32m" +} diff --git a/tests/integration/data/pretrained/minishlab___potion-base-8m_baseline.json b/tests/integration/data/pretrained/minishlab___potion-base-8m_baseline.json new file mode 100644 index 0000000..18e7133 --- /dev/null +++ b/tests/integration/data/pretrained/minishlab___potion-base-8m_baseline.json @@ -0,0 +1,73 @@ +{ + "metrics": { + "base_model_name": null, + "config": { + "apply_pca": 256, + "apply_zipf": true, + "architectures": [ + "StaticModel" + ], + "hidden_dim": 256, + "model_type": "model2vec", + "normalize": true, + "seq_length": 1000000, + "tokenizer_name": "baai/bge-base-en-v1.5" + }, + "embedding_dim": 256, + "embedding_dtype": "float32", + "embedding_mean": -0.000426, + "embedding_rank": 256, + "embedding_row_norm_mean": 94.661816, + "embedding_rows": 29528, + "embedding_std": 5.989001, + "encoding_speed": { + "sentences_per_second": 133169.33, + "tokens_per_second": 1614678.17 + }, + "first_tokens": [ + "[PAD]", + "[UNK]", + "[CLS]", + "[SEP]", + "[MASK]", + "!", + "\"", + "#", + "$", + "%" + ], + "full_vocab_size": 29528, + "has_token_mapping": false, + "has_weights": false, + "language": null, + "last_tokens": [ + "##\uff01", + "##\uff08", + "##\uff09", + "##\uff0c", + "##\uff0d", + "##\uff0e", + "##\uff0f", + "##\uff1a", + "##\uff1f", + "##\uff5e" + ], + "median_token_length": 6, + "mteb_sts_scores": { + "BIOSSES": 0.758578, + "SICK-R": 0.646753, + "STS12": 0.622476, + "STS13": 0.77276, + "STS14": 0.719141, + "STS15": 0.797535, + "STS16": 0.762634, + "STSBenchmark": 0.754052 + }, + "normalize": true, + "token_order_hash": "efd6dc3c71fd5fe1c1f97b63e7caf527b1e2196f7af18e1e1581e43e82880a8a", + "tokenizer_type": "WordPiece", + "unk_token_id": 1, + "vocabulary_quantization": null + }, + "model": "minishlab/potion-base-8m" +} diff --git a/tests/integration/data/pretrained/minishlab___potion-multilingual-128m_baseline.json b/tests/integration/data/pretrained/minishlab___potion-multilingual-128m_baseline.json new file mode 100644 index 0000000..dc5fa78 --- /dev/null +++ b/tests/integration/data/pretrained/minishlab___potion-multilingual-128m_baseline.json @@ -0,0 +1,177 @@ +{ + "metrics": { + "base_model_name": null, + "config": { + "apply_pca": null, + "apply_zipf": null, + "architectures": [ + "StaticModel" + ], + "hidden_dim": 256, + "model_type": "model2vec", + "normalize": true, + "seq_length": 1000000, + "sif_coefficient": null, + "tokenizer_name": "BAAI/bge-m3" + }, + "embedding_dim": 256, + "embedding_dtype": "float32", + "embedding_mean": -0.002286, + "embedding_rank": 256, + "embedding_row_norm_mean": 12.731157, + "embedding_rows": 500353, + "embedding_std": 0.857525, + "encoding_speed": { + "sentences_per_second": 117518.51, + "tokens_per_second": 1498360.99 + }, + "first_tokens": [ + "[PAD]", + "[UNK]", + ",", + ".", + "\u2581", + "s", + "\u2581de", + "-", + "\u2581a", + "a" + ], + "full_vocab_size": 500353, + "has_token_mapping": false, + "has_weights": false, + "language": [ + "en", + "multilingual", + "af", + "am", + "ar", + "az", + "be", + "bg", + "bn", + "ca", + "ceb", + "co", + "cs", + "cy", + "da", + "de", + "el", + "eo", + "es", + "et", + "eu", + "fa", + "fi", + "fil", + "fr", + "fy", + "ga", + "gd", + "gl", + "gu", + "ha", + "haw", + "hi", + "hmn", + "ht", + "hu", + "hy", + "id", + "ig", + "is", + "it", + "iw", + "ja", + "jv", + "ka", + "kk", + "km", + "kn", + "ko", + "ku", + "ky", + "la", + "lb", + "lo", + "lt", + "lv", + "mg", + "mi", + "mk", + "ml", + "mn", + "mr", + "ms", + "mt", + "my", + "ne", + "nl", + "no", + "ny", + "pa", + "pl", + "ps", + "pt", + "ro", + "ru", + "sd", + "si", + "sk", + "sl", + "sm", + "sn", + "so", + "sq", + "sr", + "st", + "su", + "sv", + "sw", + "ta", + "te", + "tg", + "th", + "tr", + "uk", + "ur", + "uz", + "vi", + "xh", + "yi", + "yo", + "zh", + "zu" + ], + "last_tokens": [ + "\u2581sansar", + "\u2581dprk", + "\u2581angeline", + "\u2581\u05de\u05d0\u05dc\u05d4", + "\u2581klinikos", + "\u2581dzikir", + "\u2581laivo", + "\u2581ncel", + "\u2581t38", + "\u2581lugi" + ], + "median_token_length": 7, + "mteb_sts_scores": { + "BIOSSES": 0.708733, + "SICK-R": 0.608512, + "STS12": 0.643476, + "STS13": 0.742303, + "STS14": 0.697261, + "STS15": 0.793331, + "STS16": 0.742115, + "STSBenchmark": 0.717519 + }, + "normalize": true, + "token_order_hash": "0b2eac292630557804d59ed065289d80954e3d56d793d0f7c4460fd19ec7cf22", + "tokenizer_type": "Unigram", + "unk_token_id": null, + "vocabulary_quantization": null + }, + "model": "minishlab/potion-multilingual-128m" +} diff --git a/tests/integration/data/pretrained/minishlab___potion-retrieval-32m_baseline.json b/tests/integration/data/pretrained/minishlab___potion-retrieval-32m_baseline.json new file mode 100644 index 0000000..d471192 --- /dev/null +++ b/tests/integration/data/pretrained/minishlab___potion-retrieval-32m_baseline.json @@ -0,0 +1,73 @@ +{ + "metrics": { + "base_model_name": "minishlab/potion-base-32M", + "config": { + "apply_pca": 512, + "apply_zipf": true, + "architectures": [ + "StaticModel" + ], + "hidden_dim": 512, + "model_type": "model2vec", + "normalize": true, + "seq_length": 1000000, + "tokenizer_name": "baai/bge-base-en-v1.5" + }, + "embedding_dim": 512, + "embedding_dtype": "float32", + "embedding_mean": 0.005292, + "embedding_rank": 512, + "embedding_row_norm_mean": 130.27364, + "embedding_rows": 63091, + "embedding_std": 5.912311, + "encoding_speed": { + "sentences_per_second": 120182.5, + "tokens_per_second": 1352053.17 + }, + "first_tokens": [ + "[PAD]", + "[UNK]", + "[CLS]", + "[SEP]", + "[MASK]", + "!", + "\"", + "#", + "$", + "%" + ], + "full_vocab_size": 63091, + "has_token_mapping": false, + "has_weights": false, + "language": null, + "last_tokens": [ + "nti", + "corporeal", + "downgrading", + "symbolise", + "emptive", + "daters", + "cousteau", + "crichton", + "60k", + "longines" + ], + "median_token_length": 7, + "mteb_sts_scores": { + "BIOSSES": 0.787801, + "SICK-R": 0.66508, + "STS12": 0.613768, + "STS13": 0.73606, + "STS14": 0.704635, + "STS15": 0.811203, + "STS16": 0.725174, + "STSBenchmark": 0.740959 + }, + "normalize": true, + "token_order_hash": "ba90854d713d66960327011c1d592b5f16c6dd7fe26676b103490b62753dc7c9", + "tokenizer_type": "WordPiece", + "unk_token_id": 1, + "vocabulary_quantization": null + }, + "model": "minishlab/potion-retrieval-32m" +} diff --git a/tests/integration/pretrained_model_metrics.py b/tests/integration/pretrained_model_metrics.py new file mode 100644 index 0000000..858bd6f --- /dev/null +++ b/tests/integration/pretrained_model_metrics.py @@ -0,0 +1,113 @@ +from __future__ import annotations + +import hashlib +import time +from pathlib import Path +from typing import Any + +import numpy as np + +from model2vec.model import StaticModel +from tests.integration.distill_metrics import STS_TASKS, compute_mteb_sts_scores + +PRETRAINED_MODELS: tuple[str, ...] = ( + "minishlab/potion-base-8m", + "minishlab/potion-base-32m", + "minishlab/potion-retrieval-32m", + "minishlab/potion-multilingual-128m", +) + +BASELINE_DIR = Path(__file__).parent / "data" / "pretrained" + +_SPEED_SENTENCES: tuple[str, ...] = ( + "The quick brown fox jumps over the lazy dog.", + "Paris is the capital of France, and it sits on the river Seine.", + "I would like to order a large pizza with extra cheese and mushrooms.", + "Machine learning models can be trained on very large text datasets.", + "The weather today is sunny with a light breeze coming off the sea.", + "She sells seashells by the seashore every single summer morning.", + "Static embedding models trade a little accuracy for a lot of speed.", + "hi", +) +_SPEED_CORPUS: list[str] = [_SPEED_SENTENCES[i % len(_SPEED_SENTENCES)] for i in range(2048)] +_SPEED_RUNS = 3 + + +def baseline_path_for(model_name: str) -> Path: + """The JSON baseline file for a pretrained model id, e.g. `minishlab/potion-base-8m`.""" + safe_model_name = model_name.replace("/", "___") + return BASELINE_DIR / f"{safe_model_name}_baseline.json" + + +def load_static_model(model_name: str) -> StaticModel: + """Download a published `StaticModel` from the Hugging Face hub.""" + return StaticModel.from_pretrained(model_name) + + +def measure_encoding_speed(model: StaticModel) -> dict[str, float]: + """Encode a fixed corpus a few times and report the best throughput. + + :param model: The model to benchmark. + :return: Best-of-N sentences/second and tokens/second over `_SPEED_CORPUS`. + """ + model.encode(_SPEED_CORPUS[:64], use_multiprocessing=False) + + n_tokens = sum(len(ids) for ids in model.tokenize(_SPEED_CORPUS)) + best_seconds = float("inf") + for _ in range(_SPEED_RUNS): + start = time.perf_counter() + model.encode(_SPEED_CORPUS, use_multiprocessing=False) + best_seconds = min(best_seconds, time.perf_counter() - start) + + return { + "sentences_per_second": round(len(_SPEED_CORPUS) / best_seconds, 2), + "tokens_per_second": round(n_tokens / best_seconds, 2), + } + + +def compute_metrics(model: StaticModel) -> dict[str, Any]: + """Compute a JSON-serializable snapshot of a published model's properties. + + :param model: The loaded StaticModel to summarize. + :return: A dict with every loaded attribute, embedding stats, MTEB STS scores and encoding speed. + """ + 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_dtype": model.embedding_dtype, + "embedding_rank": int(np.linalg.matrix_rank(embedding)), + "embedding_mean": round(float(embedding.mean()), 6), + "embedding_std": round(float(embedding.std()), 6), + "embedding_row_norm_mean": round(float(np.linalg.norm(embedding, axis=1).mean()), 6), + "token_order_hash": token_order_hash, + "first_tokens": tokens[:10], + "last_tokens": tokens[-10:], + "median_token_length": int(model.median_token_length), + "unk_token_id": model.unk_token_id, + "normalize": bool(model.normalize), + "base_model_name": model.base_model_name, + "language": model.language, + "vocabulary_quantization": model.vocabulary_quantization, + "has_weights": model.weights is not None, + "has_token_mapping": model.token_mapping is not None, + "tokenizer_type": type(model.tokenizer.model).__name__, + "config": model.config, + "mteb_sts_scores": compute_mteb_sts_scores(model), + "encoding_speed": measure_encoding_speed(model), + } + + +__all__ = [ + "BASELINE_DIR", + "PRETRAINED_MODELS", + "STS_TASKS", + "baseline_path_for", + "compute_metrics", + "load_static_model", + "measure_encoding_speed", +] diff --git a/tests/integration/test_pretrained_model_regression.py b/tests/integration/test_pretrained_model_regression.py new file mode 100644 index 0000000..21d5b84 --- /dev/null +++ b/tests/integration/test_pretrained_model_regression.py @@ -0,0 +1,160 @@ +from __future__ import annotations + +import json +import logging +from pathlib import Path +from typing import Any + +import numpy as np +import pytest + +logger = logging.getLogger(__name__) + +from model2vec.model import StaticModel +from tests.integration.pretrained_model_metrics import ( + PRETRAINED_MODELS, + baseline_path_for, + compute_metrics, + load_static_model, +) + +_ATTRIBUTE_FIELDS = ( + "full_vocab_size", + "embedding_rows", + "embedding_dim", + "embedding_dtype", + "token_order_hash", + "first_tokens", + "last_tokens", + "median_token_length", + "unk_token_id", + "normalize", + "base_model_name", + "language", + "vocabulary_quantization", + "has_weights", + "has_token_mapping", + "tokenizer_type", + "config", +) + + +@pytest.fixture(scope="module", params=sorted(PRETRAINED_MODELS), ids=sorted(PRETRAINED_MODELS)) +def model_name(request: pytest.FixtureRequest) -> str: + """A published model id, 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 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-pretrained-update`.") + return json.loads(path.read_text()) + + +@pytest.fixture(scope="module") +def model(model_name: str) -> StaticModel: + """Download the published model once for the whole module.""" + return load_static_model(model_name) + + +@pytest.fixture(scope="module") +def current_metrics(model: StaticModel) -> dict[str, Any]: + """Compute the same metrics as the baseline for the freshly loaded model.""" + return compute_metrics(model) + + +def test_all_attributes_are_loaded(model: StaticModel) -> None: + """Every attribute a `StaticModel` is expected to expose must be present and well-formed after loading.""" + assert isinstance(model.embedding, np.ndarray) + assert model.embedding.ndim == 2 + assert len(model.tokens) == model.embedding.shape[0] + vocab = model.tokenizer.get_vocab() + assert model.tokens == tuple(sorted(vocab, key=lambda token: vocab[token])) + assert model.dim == model.embedding.shape[1] + assert model.embedding_dtype == np.dtype(model.embedding.dtype).name + assert isinstance(model.config, dict) and model.config + assert isinstance(model.normalize, bool) + assert isinstance(model.median_token_length, int) and model.median_token_length > 0 + assert model.unk_token_id is None or isinstance(model.unk_token_id, int) + assert model.base_model_name is None or isinstance(model.base_model_name, str) + assert model.language is None or isinstance(model.language, list) + + +def test_attributes_match_baseline(model_name: str, baseline: dict[str, Any], current_metrics: dict[str, Any]) -> None: + """Vocab size, token order/identity, dtype, and every loaded attribute must match the baseline exactly.""" + expected = baseline["metrics"] + for field in _ATTRIBUTE_FIELDS: + assert current_metrics[field] == expected[field], ( + f"[{model_name}] '{field}' drifted from the baseline: " + f"expected {expected[field]!r}, got {current_metrics[field]!r}. " + "If this is intentional, run `make test-integration-pretrained-update` and review the JSON diff." + ) + + +def test_vocab_size_and_order_match_baseline( + model_name: str, baseline: dict[str, Any], current_metrics: dict[str, Any] +) -> None: + """The vocabulary must be the same size and in the same order as the baseline.""" + expected = baseline["metrics"] + assert current_metrics["full_vocab_size"] == expected["full_vocab_size"], model_name + assert current_metrics["embedding_rows"] == expected["embedding_rows"], model_name + assert current_metrics["token_order_hash"] == expected["token_order_hash"], model_name + assert current_metrics["first_tokens"] == expected["first_tokens"], model_name + assert current_metrics["last_tokens"] == expected["last_tokens"], model_name + + +def test_embedding_distribution_matches_baseline( + model_name: str, baseline: dict[str, Any], current_metrics: dict[str, Any] +) -> None: + """The embedding matrix's rank and mean/std shouldn't drift from the baseline.""" + expected = baseline["metrics"] + assert abs(current_metrics["embedding_rank"] - expected["embedding_rank"]) <= 1, model_name + assert current_metrics["embedding_mean"] == pytest.approx(expected["embedding_mean"], abs=1e-4), model_name + assert current_metrics["embedding_std"] == pytest.approx(expected["embedding_std"], rel=0.02, abs=1e-4), model_name + assert current_metrics["embedding_row_norm_mean"] == pytest.approx( + expected["embedding_row_norm_mean"], rel=0.02, abs=1e-4 + ), model_name + + +def test_mteb_sts_scores_match_baseline( + model_name: str, baseline: dict[str, Any], current_metrics: dict[str, Any] +) -> None: + """MTEB STS scores for a published model must stay within a small tolerance of the golden baseline.""" + expected_scores = baseline["metrics"]["mteb_sts_scores"] + actual_scores = current_metrics["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}] MTEB '{task_name}' STS score drifted from baseline: " + f"{expected_score:.4f} -> {actual_score:.4f}" + ) + + +def test_encoding_speed_is_measurable( + model_name: str, baseline: dict[str, Any], current_metrics: dict[str, Any] +) -> None: + """Record encoding throughput for reference. Never fails on speed: it is machine-dependent and a target to optimize.""" + expected = baseline["metrics"]["encoding_speed"] + actual = current_metrics["encoding_speed"] + logger.info( + f"[{model_name}] encoding speed: " + f"{actual['sentences_per_second']:.0f} sent/s ({expected['sentences_per_second']:.0f} baseline), " + f"{actual['tokens_per_second']:.0f} tok/s ({expected['tokens_per_second']:.0f} baseline)" + ) + assert actual["sentences_per_second"] > 0 + assert actual["tokens_per_second"] > 0 + + +def test_save_and_load_roundtrip(model: StaticModel, tmp_path: Path) -> None: + """Saving and reloading a published model must not change its tokens or embeddings.""" + save_path = tmp_path / "pretrained_model" + model.save_pretrained(save_path) + loaded_model = StaticModel.from_pretrained(save_path) + + assert loaded_model.tokens == model.tokens + assert loaded_model.unk_token_id == model.unk_token_id + assert loaded_model.normalize == model.normalize + np.testing.assert_array_equal(loaded_model.embedding, model.embedding) diff --git a/tests/integration/update_pretrained_baseline.py b/tests/integration/update_pretrained_baseline.py new file mode 100644 index 0000000..2e2ee9f --- /dev/null +++ b/tests/integration/update_pretrained_baseline.py @@ -0,0 +1,41 @@ +from __future__ import annotations + +import json +import logging +import sys + +from tests.integration.pretrained_model_metrics import ( + PRETRAINED_MODELS, + baseline_path_for, + compute_metrics, + load_static_model, +) + +logger = logging.getLogger(__name__) + + +def update_baseline(model_name: str) -> None: + """Load one published model and write its metrics snapshot to its baseline file. + + :param model_name: A model identifier on Hugging Face. + """ + model = load_static_model(model_name) + baseline = {"model": model_name, "metrics": compute_metrics(model)} + + 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 list(PRETRAINED_MODELS) + + for model_name in requested: + update_baseline(model_name) + + +if __name__ == "__main__": + logging.basicConfig(level=logging.INFO, format="%(message)s") + main() From ffd0bb830725b51d98f3e6394bffda4692c9a22d Mon Sep 17 00:00:00 2001 From: stephantul Date: Sat, 29 Aug 2026 11:13:07 +0200 Subject: [PATCH 2/2] remove all --- tests/integration/pretrained_model_metrics.py | 13 +------------ 1 file changed, 1 insertion(+), 12 deletions(-) diff --git a/tests/integration/pretrained_model_metrics.py b/tests/integration/pretrained_model_metrics.py index 858bd6f..13dce68 100644 --- a/tests/integration/pretrained_model_metrics.py +++ b/tests/integration/pretrained_model_metrics.py @@ -8,7 +8,7 @@ import numpy as np from model2vec.model import StaticModel -from tests.integration.distill_metrics import STS_TASKS, compute_mteb_sts_scores +from tests.integration.distill_metrics import compute_mteb_sts_scores PRETRAINED_MODELS: tuple[str, ...] = ( "minishlab/potion-base-8m", @@ -100,14 +100,3 @@ def compute_metrics(model: StaticModel) -> dict[str, Any]: "mteb_sts_scores": compute_mteb_sts_scores(model), "encoding_speed": measure_encoding_speed(model), } - - -__all__ = [ - "BASELINE_DIR", - "PRETRAINED_MODELS", - "STS_TASKS", - "baseline_path_for", - "compute_metrics", - "load_static_model", - "measure_encoding_speed", -]