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.
This commit is contained in:
Taras Tsugrii 2023-05-31 18:56:13 -05:00 committed by GitHub
parent bd9e0f3934
commit 0ad76c3380
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -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]: