From 205bcfdd4eef68ea01c0a1d19de9b3a84d64b7fb Mon Sep 17 00:00:00 2001 From: Ashwin Bharambe Date: Thu, 24 Oct 2024 12:11:58 -0700 Subject: [PATCH] Fix score threshold in faiss --- llama_stack/providers/impls/meta_reference/memory/faiss.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/llama_stack/providers/impls/meta_reference/memory/faiss.py b/llama_stack/providers/impls/meta_reference/memory/faiss.py index 8ead96302..02829f7be 100644 --- a/llama_stack/providers/impls/meta_reference/memory/faiss.py +++ b/llama_stack/providers/impls/meta_reference/memory/faiss.py @@ -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 )