From ebd89a0fe69a59be2b257dd3831913f44521eabc Mon Sep 17 00:00:00 2001 From: Ramzi Alashmali <186906047+RAMZI0TO99@users.noreply.github.com> Date: Sun, 4 Oct 2026 16:11:49 +0300 Subject: [PATCH 1/2] fix(sparse): preserve small positive SPLADE weights with log1p --- fastembed/sparse/splade_pp.py | 7 +- tests/test_splade_log1p.py | 131 ++++++++++++++++++++++++++++++++++ 2 files changed, 137 insertions(+), 1 deletion(-) create mode 100644 tests/test_splade_log1p.py 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_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) From 2c72313b50f3c04f46ddd1f3dbe2245584a82a23 Mon Sep 17 00:00:00 2001 From: George Panchuk Date: Mon, 5 Oct 2026 01:15:59 +0700 Subject: [PATCH 2/2] chore: drop log1p regression tests and docstring from SPLADE post-processing --- fastembed/sparse/splade_pp.py | 5 -- tests/test_splade_log1p.py | 131 ---------------------------------- 2 files changed, 136 deletions(-) delete mode 100644 tests/test_splade_log1p.py diff --git a/fastembed/sparse/splade_pp.py b/fastembed/sparse/splade_pp.py index 9dfda5e0..2b05b653 100644 --- a/fastembed/sparse/splade_pp.py +++ b/fastembed/sparse/splade_pp.py @@ -38,11 +38,6 @@ 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") diff --git a/tests/test_splade_log1p.py b/tests/test_splade_log1p.py deleted file mode 100644 index f5e02d4f..00000000 --- a/tests/test_splade_log1p.py +++ /dev/null @@ -1,131 +0,0 @@ -"""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)