Fix score threshold in faiss

This commit is contained in:
Ashwin Bharambe 2024-10-24 12:11:58 -07:00 committed by Xi Yan
parent 86f1efa680
commit 7b1a45ee0f

View file

@ -47,7 +47,9 @@ class FaissIndex(EmbeddingIndex):
self.index.add(np.array(embeddings).astype(np.float32))
async def query(self, embedding: NDArray, k: int) -> QueryDocumentsResponse:
async def query(
self, embedding: NDArray, k: int, score_threshold: float
) -> QueryDocumentsResponse:
distances, indices = self.index.search(
embedding.reshape(1, -1).astype(np.float32), k
)