Minor update to ensemble retriever to handle a mix of Documents or str (#14552)

pull/13850/head
Lance Martin 6 months ago committed by GitHub
parent ca7da8f7ef
commit 282362382c
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

@ -103,6 +103,13 @@ class EnsembleRetriever(BaseRetriever):
for i, retriever in enumerate(self.retrievers)
]
# Enforce that retrieved docs are Documents for each list in retriever_docs
for i in range(len(retriever_docs)):
retriever_docs[i] = [
Document(page_content=doc) if not isinstance(doc, Document) else doc
for doc in retriever_docs[i]
]
# apply rank fusion
fused_documents = self.weighted_reciprocal_rank(retriever_docs)
@ -130,6 +137,13 @@ class EnsembleRetriever(BaseRetriever):
for i, retriever in enumerate(self.retrievers)
]
# Enforce that retrieved docs are Documents for each list in retriever_docs
for i in range(len(retriever_docs)):
retriever_docs[i] = [
Document(page_content=doc) if not isinstance(doc, Document) else doc
for doc in retriever_docs[i]
]
# apply rank fusion
fused_documents = self.weighted_reciprocal_rank(retriever_docs)

Loading…
Cancel
Save