diff --git a/fastembed/image/transform/operators.py b/fastembed/image/transform/operators.py index c5155be4..aba1df5c 100644 --- a/fastembed/image/transform/operators.py +++ b/fastembed/image/transform/operators.py @@ -325,6 +325,11 @@ def _get_convert_to_rgb(transforms: list[Transform], config: dict[str, Any]) -> @classmethod def _get_resize(cls, transforms: list[Transform], config: dict[str, Any]) -> None: + """Append resize operations selected by the image processor configuration. + + ConvNeXT defaults to resizing when do_resize is omitted. Disabling it + skips both resize and its coupled crop while later preprocessing continues. + """ mode = config.get("image_processor_type", "CLIPImageProcessor") if mode in ("CLIPImageProcessor", "SiglipImageProcessor"): if config.get("do_resize", False): @@ -344,6 +349,8 @@ def _get_resize(cls, transforms: list[Transform], config: dict[str, Any]) -> Non ) ) elif mode == "ConvNextFeatureExtractor": + if not config.get("do_resize", True): + return if "size" in config and "shortest_edge" not in config["size"]: raise ValueError( f"Size dictionary must contain 'shortest_edge' key. Got {config['size'].keys()}" diff --git a/fastembed/parallel_processor.py b/fastembed/parallel_processor.py index 36688d87..1d39c1e4 100644 --- a/fastembed/parallel_processor.py +++ b/fastembed/parallel_processor.py @@ -112,6 +112,11 @@ def __init__( self.num_active_workers: BaseValue | None = None def start(self, **kwargs: Any) -> None: + """Start workers with copies of the supplied initialization options. + + The pool's cuda setting takes precedence over a cuda option in kwargs. + When device IDs are provided, assign them to workers in round-robin order. + """ self.emergency_shutdown = False self.input_queue = self.ctx.Queue(self.queue_size) # An emergency shutdown unblocks the feeder thread with EPIPE (see semi_ordered_map), let it @@ -125,10 +130,10 @@ def start(self, **kwargs: Any) -> None: for worker_id in range(0, self.num_workers): worker_kwargs = deepcopy(kwargs) + worker_kwargs["cuda"] = self.cuda if self.device_ids: device_id = self.device_ids[worker_id % len(self.device_ids)] worker_kwargs["device_id"] = device_id - worker_kwargs["cuda"] = self.cuda assert hasattr(self.ctx, "Process") process = self.ctx.Process( diff --git a/fastembed/sparse/bm25.py b/fastembed/sparse/bm25.py index 2b124070..ad946acd 100644 --- a/fastembed/sparse/bm25.py +++ b/fastembed/sparse/bm25.py @@ -342,6 +342,10 @@ def _term_frequency(self, tokens: list[str]) -> dict[int, float]: Returns: dict[int, float]: The token_id to term frequency mapping. + + Counts and document length refer to the processed tokens supplied here. + This experimental collision policy adds separate lexical-token weights + at a shared ID; it does not merge counts before applying the BM25 formula. """ tf_map: dict[int, float] = {} counter: defaultdict[str, int] = defaultdict(int) @@ -352,10 +356,9 @@ def _term_frequency(self, tokens: list[str]) -> dict[int, float]: for stemmed_token in counter: token_id = self.compute_token_id(stemmed_token) num_occurrences = counter[stemmed_token] - tf_map[token_id] = num_occurrences * (self.k + 1) - tf_map[token_id] /= num_occurrences + self.k * ( - 1 - self.b + self.b * doc_len / self.avg_len - ) + weight = num_occurrences * (self.k + 1) + weight /= num_occurrences + self.k * (1 - self.b + self.b * doc_len / self.avg_len) + tf_map[token_id] = tf_map.get(token_id, 0.0) + weight return tf_map @classmethod @@ -365,6 +368,9 @@ def compute_token_id(cls, token: str) -> int: def query_embed(self, query: str | Iterable[str], **kwargs: Any) -> Iterable[SparseEmbedding]: """To emulate BM25 behaviour, we don't need to use weights in the query, and it's enough to just hash the tokens and assign a weight of 1.0 to them. + + Store unique token IDs as int64 to include the absolute signed-hash + boundary value 2147483648 without changing established IDs. """ if isinstance(query, str): query = [query] @@ -375,7 +381,7 @@ def query_embed(self, query: str | Iterable[str], **kwargs: Any) -> Iterable[Spa stemmed_tokens = self._stem(tokens) token_ids = np.array( list(set(self.compute_token_id(token) for token in stemmed_tokens)), - dtype=np.int32, + dtype=np.int64, ) values = np.ones_like(token_ids) yield SparseEmbedding(indices=token_ids, values=values) diff --git a/fastembed/sparse/sparse_embedding_base.py b/fastembed/sparse/sparse_embedding_base.py index ea9a4907..2a8d781f 100644 --- a/fastembed/sparse/sparse_embedding_base.py +++ b/fastembed/sparse/sparse_embedding_base.py @@ -25,8 +25,13 @@ def as_dict(self) -> dict[int, float]: @classmethod def from_dict(cls, data: dict[int, float]) -> "SparseEmbedding": + """Build aligned indices and values from an index-to-weight mapping. + + Mapping iteration order is retained. Empty input produces empty values + and int64 indices, so the indices remain valid for NumPy indexing. + """ if len(data) == 0: - return cls(values=np.array([]), indices=np.array([])) + return cls(values=np.array([]), indices=np.array([], dtype=np.int64)) indices, values = zip(*data.items()) return cls(values=np.array(values), indices=np.array(indices)) diff --git a/fastembed/sparse/splade_pp.py b/fastembed/sparse/splade_pp.py index 1183f214..9dfda5e0 100644 --- a/fastembed/sparse/splade_pp.py +++ b/fastembed/sparse/splade_pp.py @@ -38,10 +38,15 @@ class SpladePP(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]): def _post_process_onnx_output( self, output: OnnxOutputContext, **kwargs: Any ) -> Iterable[SparseEmbedding]: + """Yield sparse scores from log-ReLU weights and masked maximum pooling. + + Use log1p to retain tiny positive logits that adding to one would lose. + An attention mask is required; dimensions with zero pooled weight are omitted. + """ if output.attention_mask is None: raise ValueError("attention_mask must be provided for document post-processing") - relu_log = np.log(1 + np.maximum(output.model_output, 0)) + relu_log = np.log1p(np.maximum(output.model_output, 0)) weighted_log = relu_log * np.expand_dims(output.attention_mask, axis=-1) diff --git a/tests/test_bm25_hash_collisions.py b/tests/test_bm25_hash_collisions.py new file mode 100644 index 00000000..82e77e38 --- /dev/null +++ b/tests/test_bm25_hash_collisions.py @@ -0,0 +1,135 @@ +"""Regression tests for order-independent BM25 weights when token hashes collide.""" + +from collections import Counter +from pathlib import Path + +import mmh3 +import pytest + +from fastembed.sparse.bm25 import Bm25 + + +COLLIDING_TOKENS = ("kitchens", "prostaglandins") +COLLIDING_TOKEN_ID = 358434922 + + +@pytest.fixture +def model() -> Bm25: + """Create only the BM25 state required to calculate term-frequency weights.""" + instance = Bm25.__new__(Bm25) + instance.k = 1.2 + instance.b = 0.75 + instance.avg_len = 256.0 + return instance + + +def _assert_bag_weights_are_order_independent(model: Bm25, tokens: list[str]) -> None: + """Check that permutations preserve hash IDs and weights within rounding tolerance.""" + expected_keys = {model.compute_token_id(token) for token in tokens} + reference = model._term_frequency(tokens) + assert set(reference) == expected_keys + for reordered in ( + list(reversed(tokens)), + sorted(tokens), + sorted(tokens, reverse=True), + tokens[1:] + tokens[:1], + ): + result = model._term_frequency(reordered) + assert set(result) == expected_keys + assert result == pytest.approx(reference, rel=1e-12, abs=1e-12) + + +def test_real_tokens_collide_after_absolute_value_of_signed_hash() -> None: + """Verify a real token pair with opposite signed hashes and the same sparse ID.""" + assert mmh3.hash("kitchens") == COLLIDING_TOKEN_ID + assert mmh3.hash("prostaglandins") == -COLLIDING_TOKEN_ID + assert {Bm25.compute_token_id(token) for token in COLLIDING_TOKENS} == {COLLIDING_TOKEN_ID} + + +@pytest.mark.parametrize("counts", [(1, 2), (3, 1), (2, 5), (7, 3)]) +def test_real_collision_weights_are_independent_of_token_order( + model: Bm25, counts: tuple[int, int] +) -> None: + """Keep unequal colliding token counts independent of first-appearance order.""" + tokens = [COLLIDING_TOKENS[0]] * counts[0] + [COLLIDING_TOKENS[1]] * counts[1] + + _assert_bag_weights_are_order_independent(model, tokens) + + +def test_real_collision_with_other_tokens_preserves_keys_and_order_independence( + model: Bm25, +) -> None: + """Preserve sparse IDs and order independence when collisions mix with other tokens.""" + tokens = ["kitchens", "hello", "kitchens", "prostaglandins", "hello", "kitchens"] + + _assert_bag_weights_are_order_independent(model, tokens) + + +@pytest.mark.parametrize("counts", [(1, 2, 3), (2, 4, 1), (5, 1, 2)]) +def test_three_distinct_colliding_tokens_are_order_independent( + model: Bm25, monkeypatch: pytest.MonkeyPatch, counts: tuple[int, int, int] +) -> None: + """Cover controlled three-way collisions without fixing a collision scoring policy.""" + colliding = ("alpha", "beta", "gamma") + original_hash = model.compute_token_id + + def controlled_hash(token: str) -> int: + """Force selected tokens to collide while preserving ordinary token hashes.""" + return COLLIDING_TOKEN_ID if token in colliding else original_hash(token) + + monkeypatch.setattr(model, "compute_token_id", controlled_hash) + tokens = [token for token, count in zip(colliding, counts) for _ in range(count)] + tokens += ["hello", "hello"] + + _assert_bag_weights_are_order_independent(model, tokens) + + +@pytest.mark.parametrize( + "tokens", + [ + ["hello"], + ["hello", "hello", "hello"], + ["hello", "world"], + ["hello", "world", "hello", "third", "world", "world"], + ], +) +def test_noncolliding_tokens_preserve_bm25_term_frequency_formula( + model: Bm25, tokens: list[str] +) -> None: + """Preserve the standard BM25 formula for ordinary and repeated noncolliding tokens.""" + counts = Counter(tokens) + token_ids = {token: model.compute_token_id(token) for token in counts} + assert len(set(token_ids.values())) == len(counts) + expected = { + token_ids[token]: count + * (model.k + 1) + / (count + model.k * (1 - model.b + model.b * len(tokens) / model.avg_len)) + for token, count in counts.items() + } + + assert model._term_frequency(tokens) == pytest.approx(expected, rel=1e-12, abs=1e-12) + + +def test_empty_tokens_have_no_term_frequencies(model: Bm25) -> None: + """Return an empty weight mapping when there are no document tokens.""" + assert model._term_frequency([]) == {} + + +def test_document_embeddings_of_the_same_colliding_token_bag_are_order_independent( + tmp_path: Path, +) -> None: + """Keep actual document embeddings stable when colliding tokens are reordered.""" + model = Bm25( + "Qdrant/bm25", + cache_dir=str(tmp_path), + specific_model_path=str(tmp_path), + disable_stemmer=True, + local_files_only=True, + ) + embeddings = list( + model.embed(["kitchens kitchens prostaglandins", "prostaglandins kitchens kitchens"]) + ) + + assert len(embeddings) == 2 + assert set(embeddings[0].as_dict()) == set(embeddings[1].as_dict()) == {COLLIDING_TOKEN_ID} + assert embeddings[0].as_dict() == pytest.approx(embeddings[1].as_dict(), rel=1e-12, abs=1e-12) diff --git a/tests/test_bm25_query_index_overflow.py b/tests/test_bm25_query_index_overflow.py new file mode 100644 index 00000000..d8df32e3 --- /dev/null +++ b/tests/test_bm25_query_index_overflow.py @@ -0,0 +1,104 @@ +"""Regression tests for BM25 query token IDs beyond the signed int32 range.""" + +from pathlib import Path + +import mmh3 +import numpy as np +import pytest + +from fastembed.sparse.bm25 import Bm25 + + +BOUNDARY_TOKEN = "ad1u66pi" +BOUNDARY_TOKEN_ID = 2**31 + + +@pytest.fixture +def model(tmp_path: Path) -> Bm25: + """Create a BM25 instance with local paths and stemming disabled.""" + return Bm25( + "Qdrant/bm25", + cache_dir=str(tmp_path), + specific_model_path=str(tmp_path), + disable_stemmer=True, + local_files_only=True, + ) + + +def test_boundary_token_hash_exceeds_signed_int32_after_absolute_value() -> None: + """Verify that a real token's absolute hash reaches the int32 overflow boundary.""" + assert mmh3.hash(BOUNDARY_TOKEN) == -(2**31) + assert Bm25.compute_token_id(BOUNDARY_TOKEN) == BOUNDARY_TOKEN_ID + + +def test_query_embedding_preserves_boundary_token_id(model: Bm25) -> None: + """Keep the boundary token's exact ID and unit query weight in int64 indices.""" + embedding = list(model.query_embed(BOUNDARY_TOKEN))[0] + + assert embedding.indices.dtype == np.int64 + assert embedding.indices.tolist() == [BOUNDARY_TOKEN_ID] + assert embedding.values.tolist() == [1] + + +def test_document_and_query_embeddings_use_same_boundary_token_id(model: Bm25) -> None: + """Use the same boundary token ID for document and query embeddings.""" + document_embedding = list(model.embed(BOUNDARY_TOKEN))[0] + query_embedding = list(model.query_embed(BOUNDARY_TOKEN))[0] + + assert ( + document_embedding.indices.tolist() + == query_embedding.indices.tolist() + == [BOUNDARY_TOKEN_ID] + ) + + +@pytest.mark.parametrize("query", ["AD1U66PI", "(ad1u66pi)!", "ad1u66pi ad1u66pi"]) +def test_query_normalization_preserves_boundary_token_id(model: Bm25, query: str) -> None: + """Preserve boundary IDs through case folding, punctuation removal, and deduplication.""" + embedding = list(model.query_embed(query))[0] + + assert embedding.indices.tolist() == [BOUNDARY_TOKEN_ID] + assert embedding.values.tolist() == [1] + + +@pytest.mark.parametrize("as_generator", [False, True]) +def test_query_iterables_handle_mixed_ordinary_and_boundary_tokens( + model: Bm25, as_generator: bool +) -> None: + """Accept lists and generators containing ordinary, boundary, and empty queries.""" + queries = ["hello", BOUNDARY_TOKEN, f"hello {BOUNDARY_TOKEN}", ""] + query_input = (query for query in queries) if as_generator else queries + + embeddings = list(model.query_embed(query_input)) + + hello_id = model.compute_token_id("hello") + expected_indices = [{hello_id}, {BOUNDARY_TOKEN_ID}, {hello_id, BOUNDARY_TOKEN_ID}, set()] + assert len(embeddings) == len(expected_indices) + for embedding, expected in zip(embeddings, expected_indices): + assert np.issubdtype(embedding.indices.dtype, np.integer) + assert set(embedding.indices.tolist()) == expected + assert embedding.values.tolist() == [1] * len(expected) + + +@pytest.mark.parametrize("query", ["hello", "hello world hello"]) +def test_ordinary_query_tokens_have_unit_weights_and_are_deduplicated( + model: Bm25, query: str +) -> None: + """Keep ordinary token hashing, deduplication, and unit query weights unchanged.""" + embedding = list(model.query_embed(query))[0] + expected_indices = {model.compute_token_id(token) for token in query.split()} + + assert np.issubdtype(embedding.indices.dtype, np.integer) + assert set(embedding.indices.tolist()) == expected_indices + assert embedding.values.tolist() == [1] * len(expected_indices) + + +def test_empty_queries_have_empty_integer_indices_and_values(model: Bm25) -> None: + """Produce empty arrays with integer indices for queries without usable tokens.""" + embeddings = list(model.query_embed(["", "!!!"])) + + assert len(embeddings) == 2 + for embedding in embeddings: + assert np.issubdtype(embedding.indices.dtype, np.integer) + assert embedding.indices.shape == (0,) + assert embedding.values.shape == (0,) diff --git a/tests/test_convnext_do_resize.py b/tests/test_convnext_do_resize.py new file mode 100644 index 00000000..d254e34c --- /dev/null +++ b/tests/test_convnext_do_resize.py @@ -0,0 +1,140 @@ +"""Regression tests for ConvNext preprocessing with resizing enabled or disabled.""" + +import numpy as np +import pytest +from PIL import Image + +from fastembed.image.transform.operators import Compose + + +def _patterned_rgb(width: int, height: int) -> Image.Image: + """Create deterministic RGB pixels that reveal unintended resizing or cropping.""" + pixels = np.random.default_rng(0).integers(0, 256, (height, width, 3), dtype=np.uint8) + return Image.fromarray(pixels) + + +@pytest.mark.parametrize("shortest_edge", [224, 384, 512]) +@pytest.mark.parametrize("image_size", [(67, 43), (43, 67), (224, 224)]) +def test_convnext_disabled_resize_preserves_shape_and_pixels( + shortest_edge: int, image_size: tuple[int, int] +) -> None: + """Preserve portrait, landscape, and square pixels when resizing is disabled.""" + image = _patterned_rgb(*image_size) + processor = Compose.from_config( + { + "image_processor_type": "ConvNextFeatureExtractor", + "do_resize": False, + "size": {"shortest_edge": shortest_edge}, + "do_rescale": False, + } + ) + + output = processor([image])[0] + + expected = np.asarray(image).transpose(2, 0, 1) + assert output.shape == expected.shape + np.testing.assert_array_equal(output, expected) + + +def test_convnext_disabled_resize_does_not_require_geometry_settings() -> None: + """Allow disabled resizing without a size dictionary or other geometry settings.""" + image = _patterned_rgb(13, 7) + processor = Compose.from_config( + { + "image_processor_type": "ConvNextFeatureExtractor", + "do_resize": False, + "do_rescale": False, + } + ) + + np.testing.assert_array_equal(processor([image])[0], np.asarray(image).transpose(2, 0, 1)) + + +@pytest.mark.parametrize("do_resize", [True, None], ids=["enabled", "omitted"]) +@pytest.mark.parametrize( + ("shortest_edge", "resize_size", "crop_box"), + [ + (224, (384, 256), (80, 16, 304, 240)), + (384, (384, 384), None), + (512, (512, 512), None), + ], +) +def test_convnext_enabled_and_default_resize_preserve_existing_behavior( + do_resize: bool | None, + shortest_edge: int, + resize_size: tuple[int, int], + crop_box: tuple[int, int, int, int] | None, +) -> None: + """Match independent Pillow results for enabled and default ConvNext geometry.""" + image = _patterned_rgb(300, 200) + config = { + "image_processor_type": "ConvNextFeatureExtractor", + "size": {"shortest_edge": shortest_edge}, + "do_rescale": False, + } + if do_resize is not None: + config["do_resize"] = do_resize + processor = Compose.from_config(config) + expected = image.resize(resize_size, Image.Resampling.BICUBIC) + if crop_box is not None: + expected = expected.crop(crop_box) + + np.testing.assert_array_equal(processor([image])[0], np.asarray(expected).transpose(2, 0, 1)) + + +@pytest.mark.parametrize("do_resize", [True, None], ids=["enabled", "omitted"]) +def test_convnext_enabled_resize_still_rejects_invalid_size(do_resize: bool | None) -> None: + """Keep shortest-edge validation active when resizing is enabled or unspecified.""" + config = { + "image_processor_type": "ConvNextFeatureExtractor", + "size": {"height": 224, "width": 224}, + } + if do_resize is not None: + config["do_resize"] = do_resize + + with pytest.raises(ValueError, match="shortest_edge"): + Compose.from_config(config) + + +def test_convnext_disabled_resize_still_rescales_and_normalizes() -> None: + """Apply rescaling and normalization even when geometry transforms are disabled.""" + image = _patterned_rgb(13, 7) + processor = Compose.from_config( + { + "image_processor_type": "ConvNextFeatureExtractor", + "do_resize": False, + "size": {"shortest_edge": 224}, + "do_rescale": True, + "rescale_factor": 0.25, + "do_normalize": True, + "image_mean": [1.0, 2.0, 3.0], + "image_std": [2.0, 4.0, 8.0], + } + ) + pixels = np.asarray(image).transpose(2, 0, 1).astype(np.float32) + expected = (pixels * 0.25 - np.array([1.0, 2.0, 3.0])[:, None, None]) / np.array( + [2.0, 4.0, 8.0] + )[:, None, None] + + np.testing.assert_array_equal(processor([image])[0], expected) + + +@pytest.mark.parametrize("mode", ["CLIPImageProcessor", "SiglipImageProcessor"]) +@pytest.mark.parametrize("do_resize", [False, True, None], ids=["disabled", "enabled", "omitted"]) +def test_other_processors_keep_existing_resize_flag_behavior( + mode: str, do_resize: bool | None +) -> None: + """Keep CLIP and Siglip's existing resize flag defaults and explicit behavior.""" + image = _patterned_rgb(67, 43) + config = { + "image_processor_type": mode, + "size": {"height": 32, "width": 48}, + "do_rescale": False, + } + if do_resize is not None: + config["do_resize"] = do_resize + expected = image.resize((48, 32), Image.Resampling.BICUBIC) if do_resize else image + + processor = Compose.from_config(config) + + np.testing.assert_array_equal(processor([image])[0], np.asarray(expected).transpose(2, 0, 1)) diff --git a/tests/test_parallel_cuda_forwarding.py b/tests/test_parallel_cuda_forwarding.py new file mode 100644 index 00000000..8fe47cf3 --- /dev/null +++ b/tests/test_parallel_cuda_forwarding.py @@ -0,0 +1,121 @@ +"""Regression tests for forwarding device choices into parallel worker startup.""" + +import os +from multiprocessing import get_all_start_methods +from typing import Any, Iterable +from unittest.mock import Mock + +import pytest + +from fastembed.common.types import Device +from fastembed.parallel_processor import ParallelWorkerPool, Worker + + +START_METHODS = [method for method in ("spawn", "forkserver") if method in get_all_start_methods()] + + +class ReportingWorker(Worker): + """Report initialization options and process IDs from real child processes.""" + + def __init__(self, options: dict[str, Any]) -> None: + """Store the initialization options received by this worker.""" + self.options = options + + @classmethod + def start(cls, **kwargs: Any) -> "ReportingWorker": + """Construct a reporter with the options passed to child startup.""" + return cls(kwargs) + + def process(self, items: Iterable[tuple[int, Any]]) -> Iterable[tuple[int, Any]]: + """Yield each indexed item with startup options and the child process ID.""" + for index, item in items: + yield index, (item, self.options, os.getpid()) + + +@pytest.mark.parametrize("start_method", START_METHODS) +@pytest.mark.parametrize("device_ids", [None, [], [2, 5]], ids=["none", "empty", "explicit"]) +@pytest.mark.parametrize("cuda", [False, True, Device.CPU, Device.CUDA, Device.AUTO]) +def test_workers_receive_cuda_selection_and_other_options( + start_method: str, device_ids: list[int] | None, cuda: bool | Device +) -> None: + """Forward CUDA choices and unrelated options with or without explicit device IDs.""" + pool = ParallelWorkerPool( + 1, ReportingWorker, start_method=start_method, device_ids=device_ids, cuda=cuda + ) + other_options = {"model_name": "reporting", "config": {"batch_size": 3}} + + results = list(pool.ordered_map(range(4), **other_options)) + + expected_options = {**other_options, "cuda": cuda} + if device_ids: + expected_options["device_id"] = device_ids[0] + assert [item for item, _, _ in results] == list(range(4)) + for _, options, pid in results: + assert options == expected_options + assert options["cuda"] is cuda + assert pid != os.getpid() + + +def test_device_ids_are_assigned_round_robin_at_worker_creation( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Check device assignment at process creation without relying on job scheduling.""" + pool = ParallelWorkerPool( + 5, ReportingWorker, start_method="spawn", device_ids=[2, 5], cuda=Device.CUDA + ) + process_factory = Mock() + process_factory.return_value.is_alive.return_value = False + monkeypatch.setattr(pool.ctx, "Process", process_factory) + other_options = {"model_name": "reporting", "config": {"batch_size": 3}} + + try: + pool.start(**other_options) + + worker_options = [call.kwargs["args"][-1] for call in process_factory.call_args_list] + assert worker_options == [ + {**other_options, "cuda": Device.CUDA, "device_id": device_id} + for device_id in [2, 5, 2, 5, 2] + ] + finally: + pool.join() + for queue in (pool.input_queue, pool.output_queue): + if queue is not None: + queue.close() + queue.join_thread() + + +@pytest.mark.parametrize("start_method", START_METHODS) +@pytest.mark.parametrize("device_ids", [None, [], [2, 5]]) +def test_pool_cuda_selection_takes_precedence_over_worker_options( + start_method: str, device_ids: list[int] | None +) -> None: + """Keep the pool's CUDA choice authoritative over conflicting worker options.""" + pool = ParallelWorkerPool( + 1, ReportingWorker, start_method=start_method, device_ids=device_ids, cuda=False + ) + + results = list(pool.ordered_map(["report"], cuda=True)) + + assert len(results) == 1 + assert results[0][1]["cuda"] is False + + +@pytest.mark.parametrize("start_method", START_METHODS) +def test_two_workers_preserve_order_and_explicit_device_options(start_method: str) -> None: + """Preserve result order and assigned device options with two real workers.""" + pool = ParallelWorkerPool( + 2, ReportingWorker, start_method=start_method, device_ids=[2, 5], cuda=False + ) + + results = list(pool.ordered_map(range(10), model_name="reporting")) + + assert [item for item, _, _ in results] == list(range(10)) + for _, options, pid in results: + assert options == { + "model_name": "reporting", + "cuda": False, + "device_id": options["device_id"], + } + assert options["device_id"] in {2, 5} + assert options["cuda"] is False + assert pid != os.getpid() diff --git a/tests/test_sparse_embedding_base.py b/tests/test_sparse_embedding_base.py new file mode 100644 index 00000000..e390cb63 --- /dev/null +++ b/tests/test_sparse_embedding_base.py @@ -0,0 +1,80 @@ +"""Regression tests for sparse embedding index types and dictionary conversion.""" + +from pathlib import Path + +import numpy as np +import pytest + +from fastembed.sparse.bm25 import Bm25 +from fastembed.sparse.sparse_embedding_base import SparseEmbedding + + +def test_from_dict_empty_indices_support_numpy_indexing() -> None: + """Keep empty embeddings usable as NumPy indices without modifying dense values.""" + embedding = SparseEmbedding.from_dict({}) + + assert embedding.indices.dtype == np.int64 + assert embedding.indices.shape == (0,) + assert embedding.values.shape == (0,) + assert embedding.as_dict() == {} + + dense = np.arange(5, dtype=float) + expected = dense.copy() + dense[embedding.indices] = embedding.values + + np.testing.assert_array_equal(dense, expected) + + +def test_from_dict_preserves_index_value_pairs_and_roundtrip() -> None: + """Preserve unsorted dictionary entries and their associated weights.""" + data = {3: 1.5, 0: 2.0, 4: -0.25} + + embedding = SparseEmbedding.from_dict(data) + + assert np.issubdtype(embedding.indices.dtype, np.integer) + np.testing.assert_array_equal(embedding.indices, [3, 0, 4]) + np.testing.assert_array_equal(embedding.values, [1.5, 2.0, -0.25]) + assert embedding.as_dict() == data + + +@pytest.mark.parametrize("empty_first", [True, False]) +def test_concatenating_empty_embedding_preserves_integer_indices(empty_first: bool) -> None: + """Retain valid integer indices when concatenating an empty embedding in either order.""" + empty = SparseEmbedding.from_dict({}) + nonempty = SparseEmbedding.from_dict({3: 1.5, 0: 2.0, 4: -0.25}) + embeddings = [empty, nonempty] if empty_first else [nonempty, empty] + + indices = np.concatenate([embedding.indices for embedding in embeddings]) + values = np.concatenate([embedding.values for embedding in embeddings]) + + assert np.issubdtype(indices.dtype, np.integer) + dense = np.zeros(5) + dense[indices] = values + + np.testing.assert_array_equal(dense, [2.0, 0.0, 0.0, 1.5, -0.25]) + + +@pytest.mark.parametrize("document", ["", "the and", "!!!"]) +def test_bm25_empty_documents_have_integer_indices(document: str, tmp_path: Path) -> None: + """Return valid empty sparse indices after BM25 removes every document token.""" + model = Bm25( + "Qdrant/bm25", + cache_dir=str(tmp_path), + specific_model_path=str(tmp_path), + disable_stemmer=True, + stopwords={"the", "and"}, + local_files_only=True, + ) + + embeddings = list(model.embed([document])) + + assert len(embeddings) == 1 + embedding = embeddings[0] + assert embedding.indices.dtype == np.int64 + assert embedding.indices.shape == (0,) + assert embedding.values.shape == (0,) + + dense = np.zeros(5) + dense[embedding.indices] = embedding.values + + np.testing.assert_array_equal(dense, np.zeros(5)) diff --git a/tests/test_splade_log1p.py b/tests/test_splade_log1p.py new file mode 100644 index 00000000..f5e02d4f --- /dev/null +++ b/tests/test_splade_log1p.py @@ -0,0 +1,131 @@ +"""Regression tests for stable SPLADE logarithms, masking, and max pooling.""" + +import math + +import numpy as np +from numpy.typing import NDArray +import pytest + +from fastembed.common.onnx_model import OnnxOutputContext +from fastembed.common.types import NumpyArray +from fastembed.sparse.sparse_embedding_base import SparseEmbedding +from fastembed.sparse.splade_pp import SpladePP + + +def _post_process( + logits: NumpyArray, attention_mask: NDArray[np.int64] | None +) -> list[SparseEmbedding]: + """Run the real SPLADE postprocessor without loading a model or performing inference.""" + model = SpladePP.__new__(SpladePP) + context = OnnxOutputContext(model_output=logits, attention_mask=attention_mask) + return list(model._post_process_onnx_output(context)) + + +@pytest.mark.parametrize( + ("dtype", "value", "relative_tolerance"), + [ + ("float16", 1e-4, 1e-3), + ("float32", 1e-8, 5e-7), + ("float64", 1e-20, 1e-14), + ("float16", 1e-3, 1e-3), + ("float32", 1e-6, 5e-7), + ("float64", 1e-12, 1e-14), + ], +) +def test_small_positive_weights_are_retained_and_accurate( + dtype: str, value: float, relative_tolerance: float +) -> None: + """Retain tiny positive logits and compare weights against a scalar log1p reference.""" + logits = np.array([[[value]]], dtype=dtype) + + embedding = _post_process(logits, np.ones((1, 1), dtype=np.int64))[0] + + assert embedding.indices.tolist() == [0] + assert embedding.values.shape == (1,) + expected = math.log1p(float(logits[0, 0, 0])) + assert embedding.values[0] == pytest.approx(expected, rel=relative_tolerance, abs=0) + + +@pytest.mark.parametrize( + ("dtype", "relative_tolerance"), + [("float16", 1e-3), ("float32", 5e-7), ("float64", 1e-14)], +) +def test_negative_and_zero_logits_remain_filtered(dtype: str, relative_tolerance: float) -> None: + """Exclude nonpositive logits while retaining ordinary positive vocabulary weights.""" + logits = np.array([[[-1.0, 0.0, 2.0], [-0.25, -0.0, 1.0]]], dtype=dtype) + + embedding = _post_process(logits, np.ones((1, 2), dtype=np.int64))[0] + + assert embedding.indices.tolist() == [2] + assert embedding.values.tolist() == pytest.approx( + [math.log1p(2.0)], rel=relative_tolerance, abs=0 + ) + + +@pytest.mark.parametrize( + ("dtype", "value", "relative_tolerance"), + [("float16", 1e-4, 1e-3), ("float32", 1e-8, 5e-7), ("float64", 1e-20, 1e-14)], +) +def test_masked_padding_does_not_replace_a_small_valid_weight( + dtype: str, value: float, relative_tolerance: float +) -> None: + """Exclude larger padded logits without dropping the valid token's tiny weight.""" + logits = np.array([[[value, 0.0], [10.0, 100.0]]], dtype=dtype) + + embedding = _post_process(logits, np.array([[1, 0]], dtype=np.int64))[0] + + assert embedding.indices.tolist() == [0] + expected = math.log1p(float(logits[0, 0, 0])) + assert embedding.values.tolist() == pytest.approx([expected], rel=relative_tolerance, abs=0) + + +@pytest.mark.parametrize("dtype", ["float16", "float32", "float64"]) +def test_all_masked_positions_produce_empty_embeddings(dtype: str) -> None: + """Return empty sparse arrays when every token position is masked out.""" + logits = np.array([[[1.0, 2.0], [3.0, 4.0]]], dtype=dtype) + + embedding = _post_process(logits, np.zeros((1, 2), dtype=np.int64))[0] + + assert embedding.indices.shape == (0,) + assert np.issubdtype(embedding.indices.dtype, np.integer) + assert embedding.values.shape == (0,) + + +def test_max_pooling_keeps_batch_order_and_ignores_padding() -> None: + """Pool only valid positions while preserving each document's place in the batch.""" + logits = np.array( + [ + [[1, 0, 2, 0], [3, 0, 1, 5], [99, 99, 99, 99]], + [[0, 7, 0, -2], [0, 4, 0, 0], [100, 100, 100, 100]], + ], + dtype=np.float32, + ) + + embeddings = _post_process(logits, np.array([[1, 1, 0], [1, 1, 0]], dtype=np.int64)) + + assert len(embeddings) == 2 + assert embeddings[0].indices.tolist() == [0, 2, 3] + assert embeddings[1].indices.tolist() == [1] + np.testing.assert_allclose( + embeddings[0].values, [math.log1p(3), math.log1p(2), math.log1p(5)], rtol=5e-7, atol=0 + ) + np.testing.assert_allclose(embeddings[1].values, [math.log1p(7)], rtol=5e-7, atol=0) + + +def test_post_processing_preserves_logits_and_attention_mask() -> None: + """Keep caller-owned logits and attention mask arrays unchanged.""" + logits = np.array([[[-2.0, 1e-8], [1.0, 3.0]]], dtype=np.float32) + attention_mask = np.array([[1, 0]], dtype=np.int64) + original_logits = logits.copy() + original_mask = attention_mask.copy() + + _post_process(logits, attention_mask) + + np.testing.assert_array_equal(logits, original_logits) + np.testing.assert_array_equal(attention_mask, original_mask) + + +def test_missing_attention_mask_keeps_existing_error() -> None: + """Retain the existing document postprocessing error for a missing attention mask.""" + with pytest.raises(ValueError, match="attention_mask must be provided"): + _post_process(np.zeros((1, 1, 2), dtype=np.float32), None)