From b61f50665e20456a165b25a98fef8e73eb7ab3d3 Mon Sep 17 00:00:00 2001 From: Taras Tsugrii Date: Wed, 31 May 2023 18:57:24 -0500 Subject: [PATCH] [retrievers][knn] Replace loop appends with list comprehension. (#5529) # Replace loop appends with list comprehension. It's much faster, more idiomatic and slightly more readable. --- langchain/retrievers/knn.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/langchain/retrievers/knn.py b/langchain/retrievers/knn.py index d6204723..362a0ec2 100644 --- a/langchain/retrievers/knn.py +++ b/langchain/retrievers/knn.py @@ -51,13 +51,14 @@ class KNNRetriever(BaseRetriever, BaseModel): denominator = np.max(similarities) - np.min(similarities) + 1e-6 normalized_similarities = (similarities - np.min(similarities)) / denominator - top_k_results = [] - for row in sorted_ix[0 : self.k]: + top_k_results = [ + Document(page_content=self.texts[row]) + for row in sorted_ix[0 : self.k] if ( self.relevancy_threshold is None or normalized_similarities[row] >= self.relevancy_threshold - ): - top_k_results.append(Document(page_content=self.texts[row])) + ) + ] return top_k_results async def aget_relevant_documents(self, query: str) -> List[Document]: