From 67dce4e96c5d4dc0620cb90feb0a8e12395c0cb8 Mon Sep 17 00:00:00 2001 From: Ramzi Alashmali <186906047+RAMZI0TO99@users.noreply.github.com> Date: Sat, 3 Oct 2026 16:23:07 +0300 Subject: [PATCH] fix: preserve integer indices in empty SparseEmbedding.from_dict results --- fastembed/sparse/sparse_embedding_base.py | 2 +- tests/test_sparse_embedding_base.py | 74 +++++++++++++++++++++++ 2 files changed, 75 insertions(+), 1 deletion(-) create mode 100644 tests/test_sparse_embedding_base.py diff --git a/fastembed/sparse/sparse_embedding_base.py b/fastembed/sparse/sparse_embedding_base.py index ea9a49075..ac497f5a7 100644 --- a/fastembed/sparse/sparse_embedding_base.py +++ b/fastembed/sparse/sparse_embedding_base.py @@ -26,7 +26,7 @@ def as_dict(self) -> dict[int, float]: @classmethod def from_dict(cls, data: dict[int, float]) -> "SparseEmbedding": 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/tests/test_sparse_embedding_base.py b/tests/test_sparse_embedding_base.py new file mode 100644 index 000000000..6b3aa9151 --- /dev/null +++ b/tests/test_sparse_embedding_base.py @@ -0,0 +1,74 @@ +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: + 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: + 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: + 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: + 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))