Skip to content

Commit f6565b9

Browse files
tk-pkm111claude
andauthored
fix(core): L2-normalize FastEmbed vectors (#843)
L2-normalizes FastEmbed output vectors at the provider boundary so SQLite vector scoring keeps its unit-vector contract for custom FastEmbed models such as multilingual MiniLM variants. Zero vectors are preserved as-is to avoid division errors, and the provider tests cover both non-unit vectors and zero-vector behavior. Verification: - uv run pytest tests/repository/test_fastembed_provider.py -q - uv run ruff check src/basic_memory/repository/fastembed_provider.py tests/repository/test_fastembed_provider.py - uv run ruff format --check src/basic_memory/repository/fastembed_provider.py tests/repository/test_fastembed_provider.py Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com> Signed-off-by: tk-pkm111 <133480534+tk-pkm111@users.noreply.github.com> Signed-off-by: phernandez <paul@basicmachines.co>
1 parent b6e8c63 commit f6565b9

2 files changed

Lines changed: 73 additions & 2 deletions

File tree

src/basic_memory/repository/fastembed_provider.py

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
from __future__ import annotations
44

55
import asyncio
6+
import math
67
from typing import TYPE_CHECKING
78

89
from loguru import logger
@@ -119,10 +120,17 @@ def _embed_batch() -> list[list[float]]:
119120
if effective_parallel is not None:
120121
embed_kwargs["parallel"] = effective_parallel
121122
vectors = list(model.embed(texts, **embed_kwargs))
123+
# sqlite_search_repository.py uses a distance-to-similarity formula that assumes
124+
# unit-normalized vectors (see the comment on line 65-67 of that file).
125+
# Some models (e.g. multilingual ones) return vectors with norm > 1, so we
126+
# L2-normalize here to satisfy that contract regardless of the chosen model.
122127
normalized: list[list[float]] = []
123128
for vector in vectors:
124-
values = vector.tolist() if hasattr(vector, "tolist") else vector
125-
normalized.append([float(value) for value in values])
129+
values = vector.tolist() if hasattr(vector, "tolist") else list(vector)
130+
norm = math.sqrt(sum(x * x for x in values))
131+
if norm > 0:
132+
values = [x / norm for x in values]
133+
normalized.append([float(v) for v in values])
126134
return normalized
127135

128136
vectors = await asyncio.to_thread(_embed_batch)

tests/repository/test_fastembed_provider.py

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
"""Tests for FastEmbedEmbeddingProvider."""
22

33
import builtins
4+
import math
45
import sys
56

67
import pytest
@@ -148,3 +149,65 @@ async def test_fastembed_provider_parallel_two_passes_multiprocessing(monkeypatc
148149
await provider.embed_documents(["parallel enabled"])
149150

150151
assert _StubTextEmbedding.last_embed_kwargs == {"batch_size": 64, "parallel": 2}
152+
153+
154+
class _UnormalizedVector:
155+
"""Stub vector with norm != 1 (simulates multilingual models like paraphrase-multilingual-*)."""
156+
157+
def __init__(self, values):
158+
self._values = values
159+
160+
def tolist(self):
161+
return self._values
162+
163+
164+
class _UnnormalizedTextEmbedding:
165+
def __init__(self, model_name: str, **_kwargs):
166+
self.model_name = model_name
167+
168+
def embed(self, texts: list[str], **_kwargs):
169+
# Return a vector with norm ~= 2.9 (typical for multilingual MiniLM models)
170+
for _ in texts:
171+
yield _UnormalizedVector([1.5, 2.0, 1.0, 0.5])
172+
173+
174+
@pytest.mark.asyncio
175+
async def test_fastembed_provider_l2_normalizes_output_vectors(monkeypatch):
176+
"""Returned vectors must be unit-normalized regardless of the raw model output.
177+
178+
sqlite_search_repository uses a formula that assumes norm == 1. Models such as
179+
paraphrase-multilingual-MiniLM-L12-v2 return vectors with norm ~2.9, which breaks
180+
cosine similarity scoring. The provider must apply L2 normalization before returning.
181+
"""
182+
module = type(sys)("fastembed")
183+
setattr(module, "TextEmbedding", _UnnormalizedTextEmbedding)
184+
monkeypatch.setitem(sys.modules, "fastembed", module)
185+
186+
provider = FastEmbedEmbeddingProvider(model_name="stub-multilingual", dimensions=4)
187+
result = await provider.embed_documents(["some text"])
188+
189+
assert len(result) == 1
190+
norm = math.sqrt(sum(x * x for x in result[0]))
191+
assert abs(norm - 1.0) < 1e-6, f"Expected unit norm, got {norm}"
192+
193+
194+
@pytest.mark.asyncio
195+
async def test_fastembed_provider_zero_vector_does_not_raise(monkeypatch):
196+
"""A zero vector from the model must be returned as-is without a division error."""
197+
198+
class _ZeroEmbedding:
199+
def __init__(self, model_name: str, **_kwargs):
200+
pass
201+
202+
def embed(self, texts: list[str], **_kwargs):
203+
for _ in texts:
204+
yield _UnormalizedVector([0.0, 0.0, 0.0, 0.0])
205+
206+
module = type(sys)("fastembed")
207+
setattr(module, "TextEmbedding", _ZeroEmbedding)
208+
monkeypatch.setitem(sys.modules, "fastembed", module)
209+
210+
provider = FastEmbedEmbeddingProvider(model_name="stub-zero", dimensions=4)
211+
result = await provider.embed_documents(["zero vector"])
212+
213+
assert result == [[0.0, 0.0, 0.0, 0.0]]

0 commit comments

Comments
 (0)