From 72f5a24fbe07e1dfd36d1db3d24ebaecd1cf70af Mon Sep 17 00:00:00 2001 From: Ramzi Alashmali <186906047+RAMZI0TO99@users.noreply.github.com> Date: Sun, 4 Oct 2026 16:11:46 +0300 Subject: [PATCH 1/2] fix: prevent BM25 query token ID overflow --- fastembed/sparse/bm25.py | 5 +- tests/test_bm25_query_index_overflow.py | 104 ++++++++++++++++++++++++ 2 files changed, 108 insertions(+), 1 deletion(-) create mode 100644 tests/test_bm25_query_index_overflow.py diff --git a/fastembed/sparse/bm25.py b/fastembed/sparse/bm25.py index 2b124070..be9fce95 100644 --- a/fastembed/sparse/bm25.py +++ b/fastembed/sparse/bm25.py @@ -365,6 +365,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 +378,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/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,) From 632c77c5efe33a1f6d532364265b22de36373f8b Mon Sep 17 00:00:00 2001 From: George Panchuk Date: Mon, 5 Oct 2026 17:02:41 +0700 Subject: [PATCH 2/2] tests: drop tests and a docstring --- fastembed/sparse/bm25.py | 3 - tests/test_bm25_query_index_overflow.py | 104 ------------------------ 2 files changed, 107 deletions(-) delete mode 100644 tests/test_bm25_query_index_overflow.py diff --git a/fastembed/sparse/bm25.py b/fastembed/sparse/bm25.py index be9fce95..e095c627 100644 --- a/fastembed/sparse/bm25.py +++ b/fastembed/sparse/bm25.py @@ -365,9 +365,6 @@ 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] diff --git a/tests/test_bm25_query_index_overflow.py b/tests/test_bm25_query_index_overflow.py deleted file mode 100644 index d8df32e3..00000000 --- a/tests/test_bm25_query_index_overflow.py +++ /dev/null @@ -1,104 +0,0 @@ -"""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,)