From 0ad76c3380055bfbd518348c3ec47548b9868d76 Mon Sep 17 00:00:00 2001 From: Taras Tsugrii Date: Wed, 31 May 2023 18:56:13 -0500 Subject: [PATCH] Replace loop appends with list comprehension. (#5528) # Replace loop appends with list comprehension. It's significantly faster because it avoids repeated method lookup. It's also more idiomatic and readable. --- langchain/retrievers/tfidf.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/langchain/retrievers/tfidf.py b/langchain/retrievers/tfidf.py index 3ccece7d79..db14f61e27 100644 --- a/langchain/retrievers/tfidf.py +++ b/langchain/retrievers/tfidf.py @@ -67,9 +67,7 @@ class TFIDFRetriever(BaseRetriever, BaseModel): results = cosine_similarity(self.tfidf_array, query_vec).reshape( (-1,) ) # Op -- (n_docs,1) -- Cosine Sim with each doc - return_docs = [] - for i in results.argsort()[-self.k :][::-1]: - return_docs.append(self.docs[i]) + return_docs = [self.docs[i] for i in results.argsort()[-self.k :][::-1]] return return_docs async def aget_relevant_documents(self, query: str) -> List[Document]: