From 17fb9960e1c1c298ba40759d8544deb28bb74f24 Mon Sep 17 00:00:00 2001
From: yunyv <62236062+yunyv@users.noreply.github.com>
Date: Wed, 28 Oct 2026 02:55:38 -0600
Subject: [PATCH] fix(search): keep zero float16 cosine scores finite (#674)

* fix(search): keep zero float16 cosine scores finite

Updated the handling of zero-norm rows to use a representable denominator of 1 instead of 1e-10.

* test(search): cover zero cosine vectors and batch boundaries

Added tests for batch cosine similarity with zero vectors and empty batches.

* test(search): keep zero embedding rankings finite

Add tests for zero embeddings in vector search and legacy JSON loading.
---
 reme/utils/similarity_utils.py            |  4 +--
 tests/unit/test_file_store_consistency.py | 21 +++++++++++++
 tests/unit/test_utils.py                  | 36 +++++++++++++++++++++++
 3 files changed, 59 insertions(+), 2 deletions(-)

diff --git a/reme/utils/similarity_utils.py b/reme/utils/similarity_utils.py
index ab5d7bba..0452a0e1 100644
--- a/reme/utils/similarity_utils.py
+++ b/reme/utils/similarity_utils.py
@@ -33,7 +33,7 @@ def batch_cosine_similarity(nd_array1: np.ndarray, nd_array2: np.ndarray) -> np.
     norms1 = np.linalg.norm(nd_array1, axis=1)
     norms2 = np.linalg.norm(nd_array2, axis=1)
     norm_products = np.outer(norms1, norms2)
-    # Guard against zero-norm rows to keep division finite.
-    norm_products = np.where(norm_products == 0, 1e-10, norm_products)
+    # Use a representable denominator even in float16; zero rows have zero dot products.
+    norm_products = np.where(norm_products == 0, 1, norm_products)
 
     return dot_products / norm_products
diff --git a/tests/unit/test_file_store_consistency.py b/tests/unit/test_file_store_consistency.py
index 346a507e..738db1b1 100644
--- a/tests/unit/test_file_store_consistency.py
+++ b/tests/unit/test_file_store_consistency.py
@@ -522,6 +522,27 @@ async def test_vector_search_batches_candidates_and_preserves_stable_ties(monkey
     assert [result.id for result in results] == ["0", "1", "2"]
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("zero_query", [False, True])
+async def test_vector_search_zero_embeddings_have_finite_scores(zero_query):
+    """Zero float16 embeddings do not leak NaN scores into local search ranking."""
+    store = LocalFileStore(name="t_vector_zero", embedding_store="")
+    store.embedding_store = FakeEmbeddingStore()
+    query = np.array([0.0, 0.0] if zero_query else [1.0, 0.0], dtype=np.float16)
+    store.embedding_store.get_embedding = AsyncMock(return_value=query)
+    for chunk_id, embedding in [("zero", [0.0, 0.0]), ("match", [1.0, 0.0])]:
+        candidate = chunk(chunk_id, f"{chunk_id}.md", "alpha")
+        candidate.embedding = np.array(embedding, dtype=np.float16)
+        store.file_chunks[candidate.id] = candidate
+
+    with np.errstate(invalid="raise", divide="raise"):
+        results = await store.vector_search("alpha", 2, {})
+
+    expected_ids = ["zero", "match"] if zero_query else ["match", "zero"]
+    assert [result.id for result in results] == expected_ids
+    assert [result.scores["vector"] for result in results] == ([0.0, 0.0] if zero_query else [1.0, 0.0])
+
+
 @pytest.mark.asyncio
 async def test_chunk_persistence_loads_legacy_json_embedding_list():
     """Existing indexes with JSON float-list embeddings remain readable."""
diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py
index 6b21eef1..756fecc7 100644
--- a/tests/unit/test_utils.py
+++ b/tests/unit/test_utils.py
@@ -38,6 +38,42 @@ def test_batch_cosine_similarity_pairwise_matrix():
     np.testing.assert_allclose(result[1], [0.0, 2**-0.5])
 
 
+@pytest.mark.parametrize("dtype", [np.float16, np.float32, np.float64])
+def test_batch_cosine_similarity_zero_vectors(dtype):
+    """Zero rows and columns have finite zero scores, including stored float16 embeddings."""
+    left = np.array([[0, 0], [1, 0], [0, 1]], dtype=dtype)
+    right = np.array([[1, 0], [0, 0], [-1, 0]], dtype=dtype)
+
+    with np.errstate(invalid="raise", divide="raise"):
+        result = batch_cosine_similarity(left, right)
+        reverse = batch_cosine_similarity(right, left)
+        zeros = batch_cosine_similarity(np.zeros((2, 2), dtype=dtype), right)
+        all_zeros = batch_cosine_similarity(np.zeros((2, 2), dtype=dtype), np.zeros((3, 2), dtype=dtype))
+
+    assert result.dtype == dtype
+    np.testing.assert_array_equal(result, [[0, 0, 0], [1, 0, -1], [0, 0, 0]])
+    np.testing.assert_array_equal(reverse, result.T)
+    np.testing.assert_array_equal(zeros, np.zeros((2, 3)))
+    np.testing.assert_array_equal(all_zeros, np.zeros((2, 3)))
+
+
+@pytest.mark.parametrize("left_rows,right_rows", [(0, 2), (2, 0), (0, 0)])
+def test_batch_cosine_similarity_empty_batches(left_rows, right_rows):
+    """Empty batches retain the pairwise output shape."""
+    result = batch_cosine_similarity(
+        np.zeros((left_rows, 2), dtype=np.float16),
+        np.zeros((right_rows, 2), dtype=np.float16),
+    )
+    assert result.shape == (left_rows, right_rows)
+    assert result.dtype == np.float16
+
+
+def test_batch_cosine_similarity_rejects_mismatched_dimensions():
+    """Batched embeddings must share their feature dimension."""
+    with pytest.raises(ValueError, match="Embedding dimensions must match"):
+        batch_cosine_similarity(np.zeros((2, 2)), np.zeros((2, 3)))
+
+
 def test_cosine_similarity_rejects_mismatched_lengths():
     """Single-vector cosine validates dimensions before computing."""
     with pytest.raises(ValueError, match="Vectors must have same length"):