fix(faiss): normalize vectors for cosine distance strategy (#5960)
This commit is contained in:
@@ -248,6 +248,19 @@ class FAISS(VectorStoreBase):
|
|||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.warning(f"Failed to save FAISS index: {e}")
|
logger.warning(f"Failed to save FAISS index: {e}")
|
||||||
|
|
||||||
|
def _should_normalize(self) -> bool:
|
||||||
|
"""Whether vectors must be L2-normalized before indexing/searching.
|
||||||
|
|
||||||
|
Cosine similarity is implemented on top of an inner-product index
|
||||||
|
(``IndexFlatIP``), which only equals cosine when the inputs are unit
|
||||||
|
vectors — so cosine *always* requires normalization. For euclidean the
|
||||||
|
``normalize_L2`` flag remains an opt-in.
|
||||||
|
"""
|
||||||
|
strategy = self.distance_strategy.lower()
|
||||||
|
if strategy == "cosine":
|
||||||
|
return True
|
||||||
|
return self.normalize_L2 and strategy == "euclidean"
|
||||||
|
|
||||||
def _parse_output(self, scores, ids, top_k=None) -> List[OutputData]:
|
def _parse_output(self, scores, ids, top_k=None) -> List[OutputData]:
|
||||||
"""
|
"""
|
||||||
Parse the output data.
|
Parse the output data.
|
||||||
@@ -347,7 +360,7 @@ class FAISS(VectorStoreBase):
|
|||||||
|
|
||||||
vectors_np = np.array(vectors, dtype=np.float32)
|
vectors_np = np.array(vectors, dtype=np.float32)
|
||||||
|
|
||||||
if self.normalize_L2 and self.distance_strategy.lower() == "euclidean":
|
if self._should_normalize():
|
||||||
faiss.normalize_L2(vectors_np)
|
faiss.normalize_L2(vectors_np)
|
||||||
|
|
||||||
self.index.add(vectors_np)
|
self.index.add(vectors_np)
|
||||||
@@ -384,7 +397,7 @@ class FAISS(VectorStoreBase):
|
|||||||
if len(query_vectors.shape) == 1:
|
if len(query_vectors.shape) == 1:
|
||||||
query_vectors = query_vectors.reshape(1, -1)
|
query_vectors = query_vectors.reshape(1, -1)
|
||||||
|
|
||||||
if self.normalize_L2 and self.distance_strategy.lower() == "euclidean":
|
if self._should_normalize():
|
||||||
faiss.normalize_L2(query_vectors)
|
faiss.normalize_L2(query_vectors)
|
||||||
|
|
||||||
fetch_k = top_k * 2 if filters else top_k
|
fetch_k = top_k * 2 if filters else top_k
|
||||||
|
|||||||
@@ -704,3 +704,53 @@ class TestFAISSSecurityIntegration:
|
|||||||
assert not os.path.exists(json_path), "JSON file should be deleted"
|
assert not os.path.exists(json_path), "JSON file should be deleted"
|
||||||
assert not os.path.exists(pkl_path), "PKL file should be deleted"
|
assert not os.path.exists(pkl_path), "PKL file should be deleted"
|
||||||
assert not os.path.exists(faiss_index_path), "FAISS index should be deleted"
|
assert not os.path.exists(faiss_index_path), "FAISS index should be deleted"
|
||||||
|
|
||||||
|
|
||||||
|
class TestCosineNormalization:
|
||||||
|
"""Cosine distance must rank by angle, not raw inner-product magnitude.
|
||||||
|
|
||||||
|
Regression test for the bug where cosine used an IndexFlatIP index but never
|
||||||
|
L2-normalized vectors, so results were ranked by inner product instead of
|
||||||
|
cosine similarity.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def test_cosine_ranks_by_angle_not_magnitude(self):
|
||||||
|
# Query is perfectly aligned with A (cosine 1.0) but A has a small
|
||||||
|
# magnitude, so its inner product (0.1) is lower than B's (0.5).
|
||||||
|
# Under correct cosine ranking, A must come first regardless.
|
||||||
|
with tempfile.TemporaryDirectory() as temp_dir:
|
||||||
|
store = FAISS(
|
||||||
|
collection_name="cosine_col",
|
||||||
|
path=os.path.join(temp_dir, "cosine"),
|
||||||
|
distance_strategy="cosine",
|
||||||
|
embedding_model_dims=2,
|
||||||
|
)
|
||||||
|
store.insert(
|
||||||
|
vectors=[[0.1, 0.0], [0.5, 0.5]],
|
||||||
|
payloads=[{"name": "A"}, {"name": "B"}],
|
||||||
|
ids=["A", "B"],
|
||||||
|
)
|
||||||
|
|
||||||
|
results = store.search(query="", vectors=[1.0, 0.0], top_k=2)
|
||||||
|
|
||||||
|
assert [r.id for r in results] == ["A", "B"]
|
||||||
|
# Scores are true cosine similarities, not raw inner products.
|
||||||
|
assert results[0].score == pytest.approx(1.0, abs=1e-5)
|
||||||
|
assert results[1].score == pytest.approx(0.70710677, abs=1e-5)
|
||||||
|
|
||||||
|
def test_cosine_normalizes_on_insert_and_search(self):
|
||||||
|
# A non-unit query that points the same direction as a stored vector
|
||||||
|
# should score ~1.0 once both sides are normalized.
|
||||||
|
with tempfile.TemporaryDirectory() as temp_dir:
|
||||||
|
store = FAISS(
|
||||||
|
collection_name="cosine_col2",
|
||||||
|
path=os.path.join(temp_dir, "cosine2"),
|
||||||
|
distance_strategy="cosine",
|
||||||
|
embedding_model_dims=3,
|
||||||
|
)
|
||||||
|
store.insert(vectors=[[3.0, 0.0, 0.0]], payloads=[{"name": "x"}], ids=["x"])
|
||||||
|
|
||||||
|
results = store.search(query="", vectors=[7.0, 0.0, 0.0], top_k=1)
|
||||||
|
|
||||||
|
assert results[0].id == "x"
|
||||||
|
assert results[0].score == pytest.approx(1.0, abs=1e-5)
|
||||||
|
|||||||
Reference in New Issue
Block a user