TBNilles/LEANN
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: Improve OpenAI embeddings handling in HNSW backend

Browse Source
This commit is contained in:
Andy Lee
2025-07-28 10:15:56 -07:00
parent 41812c7d22
commit 34ef0db42f
3 changed files with 23 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
packages/leann-backend-hnsw/leann_backend_hnsw/hnsw_backend.py
View File

@@ -124,7 +124,9 @@ class HNSWSearcher(BaseSearcher):
) )
from . import faiss # type: ignore from . import faiss # type: ignore
self.distance_metric = self.meta.get("distance_metric", "mips").lower() self.distance_metric = (
self.meta.get("backend_kwargs", {}).get("distance_metric", "mips").lower()
)
metric_enum = get_metric_map().get(self.distance_metric) metric_enum = get_metric_map().get(self.distance_metric)
if metric_enum is None: if metric_enum is None:
raise ValueError(f"Unsupported distance_metric '{self.distance_metric}'.") raise ValueError(f"Unsupported distance_metric '{self.distance_metric}'.")
@@ -200,6 +202,16 @@ class HNSWSearcher(BaseSearcher):
params.efSearch = complexity params.efSearch = complexity
params.beam_size = beam_width params.beam_size = beam_width
# For OpenAI embeddings with cosine distance, disable relative distance check
# This prevents early termination when all scores are in a narrow range
embedding_model = self.meta.get("embedding_model", "").lower()
if self.distance_metric == "cosine" and any(
openai_model in embedding_model for openai_model in ["text-embedding", "openai"]
):
params.check_relative_distance = False
else:
params.check_relative_distance = True
# PQ pruning: direct mapping to HNSW's pq_pruning_ratio # PQ pruning: direct mapping to HNSW's pq_pruning_ratio
params.pq_pruning_ratio = prune_ratio params.pq_pruning_ratio = prune_ratio

2
packages/leann-core/src/leann/embedding_server_manager.py
View File

@@ -293,6 +293,8 @@ class EmbeddingServerManager:
command.extend(["--passages-file", str(passages_file)]) command.extend(["--passages-file", str(passages_file)])
if embedding_mode != "sentence-transformers": if embedding_mode != "sentence-transformers":
command.extend(["--embedding-mode", embedding_mode]) command.extend(["--embedding-mode", embedding_mode])
if kwargs.get("distance_metric"):
command.extend(["--distance-metric", kwargs["distance_metric"]])
return command return command

9
packages/leann-core/src/leann/searcher_base.py
View File

@@ -63,12 +63,19 @@ class BaseSearcher(LeannBackendSearcherInterface, ABC):
if not self.embedding_model: if not self.embedding_model:
raise ValueError("Cannot use recompute mode without 'embedding_model' in meta.json.") raise ValueError("Cannot use recompute mode without 'embedding_model' in meta.json.")
# Get distance_metric from meta if not provided in kwargs
distance_metric = (
kwargs.get("distance_metric")
or self.meta.get("backend_kwargs", {}).get("distance_metric")
or "mips"
)
server_started, actual_port = self.embedding_server_manager.start_server( server_started, actual_port = self.embedding_server_manager.start_server(
port=port, port=port,
model_name=self.embedding_model, model_name=self.embedding_model,
embedding_mode=self.embedding_mode, embedding_mode=self.embedding_mode,
passages_file=passages_source_file, passages_file=passages_source_file,
distance_metric=kwargs.get("distance_metric"), distance_metric=distance_metric,
enable_warmup=kwargs.get("enable_warmup", False), enable_warmup=kwargs.get("enable_warmup", False),
) )
if not server_started: if not server_started: