diff --git a/fastembed/text/custom_text_embedding.py b/fastembed/text/custom_text_embedding.py index 512b914bd..a4c49fee0 100644 --- a/fastembed/text/custom_text_embedding.py +++ b/fastembed/text/custom_text_embedding.py @@ -13,6 +13,7 @@ from fastembed.common.types import NumpyArray, Device from fastembed.common.utils import normalize, mean_pooling, last_token_pooling from fastembed.text.onnx_embedding import OnnxTextEmbedding +from fastembed.text.onnx_text_model import TextEmbeddingWorker @dataclass(frozen=True) @@ -58,6 +59,27 @@ def __init__( def _list_supported_models(cls) -> list[DenseModelDescription]: return cls.SUPPORTED_MODELS + def embed( + self, + documents: str | Iterable[str], + batch_size: int = 256, + parallel: int | None = None, + **kwargs: Any, + ) -> Iterable[NumpyArray]: + # Spawned workers cannot see models registered in the parent process. + yield from super().embed( + documents, + batch_size, + parallel, + _custom_model_description=self.model_description, + _custom_postprocessing_config=PostprocessingConfig(self._pooling, self._normalization), + **kwargs, + ) + + @classmethod + def _get_worker_class(cls) -> type[TextEmbeddingWorker[NumpyArray]]: + return CustomTextEmbeddingWorker + def _post_process_onnx_output( self, output: OnnxOutputContext, **kwargs: Any ) -> Iterable[NumpyArray]: @@ -102,3 +124,25 @@ def add_model( cls.POSTPROCESSING_MAPPING[model_description.model] = PostprocessingConfig( pooling=pooling, normalization=normalization ) + + +class CustomTextEmbeddingWorker(TextEmbeddingWorker[NumpyArray]): + def init_embedding( + self, + model_name: str, + cache_dir: str, + **kwargs: Any, + ) -> CustomTextEmbedding: + model_description: DenseModelDescription = kwargs.pop("_custom_model_description") + postprocessing_config: PostprocessingConfig = kwargs.pop("_custom_postprocessing_config") + CustomTextEmbedding.add_model( + model_description, + pooling=postprocessing_config.pooling, + normalization=postprocessing_config.normalization, + ) + return CustomTextEmbedding( + model_name=model_name, + cache_dir=cache_dir, + threads=1, + **kwargs, + ) diff --git a/tests/misc/generate_token_lookup.py b/tests/misc/generate_token_lookup.py new file mode 100644 index 000000000..22b9244fd --- /dev/null +++ b/tests/misc/generate_token_lookup.py @@ -0,0 +1,19 @@ +"""Regenerate token_lookup.onnx; requires the development-only onnx package.""" + +from pathlib import Path + +import onnx +from onnx import TensorProto, helper + +table = helper.make_tensor("table", TensorProto.FLOAT, [5, 2], [0, 0, 0, 0, 2, 4, 3, 9, 4, 16]) +graph = helper.make_graph( + [helper.make_node("Gather", ["table", "input_ids"], ["embeddings"], axis=0)], + "synthetic-token-lookup", + [helper.make_tensor_value_info("input_ids", TensorProto.INT64, ["batch", "sequence"])], + [helper.make_tensor_value_info("embeddings", TensorProto.FLOAT, ["batch", "sequence", 2])], + [table], +) +model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)]) +model.ir_version = 8 +onnx.checker.check_model(model) +onnx.save(model, Path(__file__).with_name("token_lookup.onnx")) diff --git a/tests/misc/token_lookup.md b/tests/misc/token_lookup.md new file mode 100644 index 000000000..a52ad84b3 --- /dev/null +++ b/tests/misc/token_lookup.md @@ -0,0 +1,36 @@ +# Synthetic token lookup fixture + +`token_lookup.onnx` is an original synthetic graph for offline process-parallel +embedding tests. It contains no trained weights or external model data. + +The graph is one `Gather` node (axis 0), with an int64 `input_ids` input of shape +`[batch, sequence]` and float32 output `[batch, sequence, 2]`. Its constant table is: + +| Token ID | Vector | +| --- | --- | +| 0 (unknown) | [0, 0] | +| 1 (padding) | [0, 0] | +| 2 (alpha) | [2, 4] | +| 3 (beta) | [3, 9] | +| 4 (gamma) | [4, 16] | + +This makes CLS, mean and last-token pooling expectations directly calculable. +Tests generate their tokenizer metadata in a temporary directory and run actual +ONNX Runtime inference. The fixture uses ONNX opset 13 and IR version 8. + +Regenerate with the development dependencies installed: + +```sh +python tests/misc/generate_token_lookup.py +``` + +Generated with ONNX 1.23.0; size 234 bytes; SHA-256: + +```text +8ff79ae726a3cb5cad3c6b4f6c458dc2e085f815c13f6117ba4df106911c26bf +``` + +The generator also runs `onnx.checker.check_model`. Normal pytest execution loads +the checked-in file through ONNX Runtime and does not import the development-only +`onnx` package. The builtin-identity control deliberately uses this synthetic +graph too; it tests dispatch/inference parity, not pretrained model quality. diff --git a/tests/misc/token_lookup.onnx b/tests/misc/token_lookup.onnx new file mode 100644 index 000000000..747cfa617 Binary files /dev/null and b/tests/misc/token_lookup.onnx differ diff --git a/tests/test_custom_model_parallel.py b/tests/test_custom_model_parallel.py new file mode 100644 index 000000000..fc972e8b3 --- /dev/null +++ b/tests/test_custom_model_parallel.py @@ -0,0 +1,154 @@ +"""Public custom-model embedding must preserve serial behavior in worker processes.""" + +import json +import multiprocessing +import shutil +from pathlib import Path + +import numpy as np +import pytest +from tokenizers import Tokenizer +from tokenizers.models import WordLevel +from tokenizers.pre_tokenizers import Whitespace + +from fastembed import TextEmbedding +from fastembed.common.model_description import ModelSource, PoolingType +from fastembed.text import onnx_text_model +from fastembed.text.custom_text_embedding import CustomTextEmbedding + +START_METHODS = [ + method + for method in ("spawn", "forkserver") + if method in multiprocessing.get_all_start_methods() +] +DOCUMENTS = ["alpha beta", "gamma", "beta alpha gamma", "alpha", "gamma beta", "beta"] +EXPECTED = { + PoolingType.MEAN: [[2.5, 6.5], [4, 16], [3, 29 / 3], [2, 4], [3.5, 12.5], [3, 9]], + PoolingType.CLS: [[2, 4], [4, 16], [3, 9], [2, 4], [4, 16], [3, 9]], + PoolingType.LAST_TOKEN: [[3, 9], [4, 16], [4, 16], [2, 4], [3, 9], [3, 9]], +} + + +@pytest.fixture +def local_model(tmp_path, monkeypatch): + monkeypatch.setattr(CustomTextEmbedding, "SUPPORTED_MODELS", []) + monkeypatch.setattr(CustomTextEmbedding, "POSTPROCESSING_MAPPING", {}) + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + fixture = Path(__file__).parent / "misc" / "token_lookup.onnx" + shutil.copyfile(fixture, tmp_path / "custom-lookup.onnx") + tokenizer = Tokenizer( + WordLevel({"[UNK]": 0, "[PAD]": 1, "alpha": 2, "beta": 3, "gamma": 4}, unk_token="[UNK]") + ) + tokenizer.pre_tokenizer = Whitespace() + tokenizer.save(str(tmp_path / "tokenizer.json")) + for filename, metadata in { + "config.json": {"pad_token_id": 1}, + "tokenizer_config.json": {"model_max_length": 32, "pad_token": "[PAD]"}, + "special_tokens_map.json": {"unk_token": "[UNK]", "pad_token": "[PAD]"}, + }.items(): + (tmp_path / filename).write_text(json.dumps(metadata), encoding="utf-8") + return tmp_path + + +def register_model(name, pooling, normalization): + TextEmbedding.add_custom_model( + name, + pooling=pooling, + normalization=normalization, + sources=ModelSource(hf="synthetic-unused/local-model"), + dim=2, + model_file="custom-lookup.onnx", + ) + + +def make_model(directory, name, lazy_load): + return TextEmbedding( + name, + cache_dir=str(directory / "cache"), + specific_model_path=str(directory), + local_files_only=True, + providers=["CPUExecutionProvider"], + threads=1, + lazy_load=lazy_load, + ) + + +def expected_vectors(pooling, normalization): + values = np.array(EXPECTED[pooling], dtype=np.float32) + if normalization: + values = values / np.linalg.norm(values, axis=1, keepdims=True) + return values + + +@pytest.mark.parametrize("start_method", START_METHODS) +@pytest.mark.parametrize( + "pooling,normalization,lazy_load,as_generator,batch_size", + [ + (PoolingType.MEAN, False, True, True, 2), + (PoolingType.LAST_TOKEN, True, False, False, 1), + (PoolingType.CLS, False, True, False, 1), + ], +) +def test_custom_model_parallel_matches_serial( + local_model, + monkeypatch, + start_method, + pooling, + normalization, + lazy_load, + as_generator, + batch_size, +): + # Select a real supported OS process mode; never replace the pool or worker implementation. + monkeypatch.setattr(onnx_text_model, "get_all_start_methods", lambda: [start_method]) + register_model("Synthetic/Other", PoolingType.CLS, True) + register_model("Synthetic/Target", pooling, normalization) + model = make_model(local_model, "synthetic/target", lazy_load) + expected = expected_vectors(pooling, normalization) + serial = np.array(list(model.embed(DOCUMENTS, batch_size=batch_size))) + np.testing.assert_allclose(serial, expected, rtol=1e-6, atol=1e-7) + documents = (text for text in DOCUMENTS) if as_generator else DOCUMENTS + parallel = np.array(list(model.embed(documents, batch_size=batch_size, parallel=2))) + np.testing.assert_allclose(parallel, expected, rtol=1e-6, atol=1e-7) + + # A separate registered configuration must not inherit target pooling/normalization. + other = make_model(local_model, "Synthetic/Other", True) + other_parallel = np.array(list(other.embed(DOCUMENTS, batch_size=2, parallel=2))) + np.testing.assert_allclose( + other_parallel, expected_vectors(PoolingType.CLS, True), rtol=1e-6, atol=1e-7 + ) + + +@pytest.mark.parametrize("start_method", START_METHODS) +def test_builtin_identity_parallel_control(local_model, monkeypatch, start_method): + monkeypatch.setattr(onnx_text_model, "get_all_start_methods", lambda: [start_method]) + # The same synthetic graph checks worker infrastructure, not pretrained model quality. + shutil.copyfile(local_model / "custom-lookup.onnx", local_model / "model_optimized.onnx") + model = make_model(local_model, "BAAI/bge-small-en-v1.5", True) + serial = np.array(list(model.embed(DOCUMENTS, batch_size=2))) + parallel = np.array(list(model.embed(DOCUMENTS, batch_size=2, parallel=2))) + np.testing.assert_allclose( + serial, expected_vectors(PoolingType.CLS, True), rtol=1e-6, atol=1e-7 + ) + np.testing.assert_allclose(parallel, serial, rtol=1e-6, atol=1e-7) + + +def test_custom_small_list_parallel_flag_control(local_model): + register_model("Synthetic/Small", PoolingType.MEAN, False) + model = make_model(local_model, "Synthetic/Small", True) + result = np.array(list(model.embed(DOCUMENTS, batch_size=len(DOCUMENTS) + 1, parallel=2))) + np.testing.assert_allclose( + result, expected_vectors(PoolingType.MEAN, False), rtol=1e-6, atol=1e-7 + ) + + +@pytest.mark.parametrize("start_method", START_METHODS) +def test_custom_lazy_model_parallel_first(local_model, monkeypatch, start_method): + monkeypatch.setattr(onnx_text_model, "get_all_start_methods", lambda: [start_method]) + register_model("Synthetic/LazyFirst", PoolingType.MEAN, True) + model = make_model(local_model, "synthetic/lazyfirst", True) + # No serial inference or explicit load precedes the first public embedding call. + result = np.array(list(model.embed(DOCUMENTS, batch_size=2, parallel=2))) + np.testing.assert_allclose( + result, expected_vectors(PoolingType.MEAN, True), rtol=1e-6, atol=1e-7 + )