diff --git a/docs/docs/integrations/retrievers/flashrank-reranker.ipynb b/docs/docs/integrations/retrievers/flashrank-reranker.ipynb index 93697614cd7..bdd4ed6d762 100644 --- a/docs/docs/integrations/retrievers/flashrank-reranker.ipynb +++ b/docs/docs/integrations/retrievers/flashrank-reranker.ipynb @@ -53,7 +53,10 @@ "def pretty_print_docs(docs):\n", " print(\n", " f\"\\n{'-' * 100}\\n\".join(\n", - " [f\"Document {i+1}:\\n\\n\" + d.page_content for i, d in enumerate(docs)]\n", + " [\n", + " f\"Document {i+1}:\\n\\n{d.page_content}\\nMetadata: {d.metadata}\"\n", + " for i, d in enumerate(docs)\n", + " ]\n", " )\n", " )" ] @@ -316,6 +319,8 @@ ").load()\n", "text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=100)\n", "texts = text_splitter.split_documents(documents)\n", + "for idx, text in enumerate(texts):\n", + " text.metadata[\"id\"] = idx\n", "\n", "embedding = OpenAIEmbeddings(model=\"text-embedding-ada-002\")\n", "retriever = FAISS.from_documents(texts, embedding).as_retriever(search_kwargs={\"k\": 20})\n", diff --git a/libs/langchain/langchain/retrievers/document_compressors/flashrank_rerank.py b/libs/langchain/langchain/retrievers/document_compressors/flashrank_rerank.py index 273aff3d44c..a05fd8c6e1c 100644 --- a/libs/langchain/langchain/retrievers/document_compressors/flashrank_rerank.py +++ b/libs/langchain/langchain/retrievers/document_compressors/flashrank_rerank.py @@ -59,16 +59,20 @@ class FlashrankRerank(BaseDocumentCompressor): callbacks: Optional[Callbacks] = None, ) -> Sequence[Document]: passages = [ - {"id": i, "text": doc.page_content} for i, doc in enumerate(documents) + {"id": i, "text": doc.page_content, "meta": doc.metadata} + for i, doc in enumerate(documents) ] rerank_request = RerankRequest(query=query, passages=passages) rerank_response = self.client.rerank(rerank_request)[: self.top_n] final_results = [] + for r in rerank_response: + metadata = r["meta"] + metadata["relevance_score"] = r["score"] doc = Document( page_content=r["text"], - metadata={"id": r["id"], "relevance_score": r["score"]}, + metadata=metadata, ) final_results.append(doc) return final_results