github/langchain
1
0
Fork 0
mirror of https://github.com/hwchase17/langchain.git synced 2025-06-01 04:29:09 +00:00
Code Issues Packages Projects Releases Wiki Activity

cohere[patch]: Fix cohere rerank (#19624)

Browse Source 
Fix cohere rerank inspired by
https://github.com/langchain-ai/langchain/pull/19486
This commit is contained in:
billytrend-cohere 2024-03-27 10:41:53 -05:00 committed by GitHub
parent 8ab7bb3166
commit 85f57ab4cd
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 22 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
libs/partners/cohere/langchain_cohere/rerank.py
View File

@ -69,10 +69,14 @@ class CohereRerank(BaseDocumentCompressor):
model = model or self.model
top_n = top_n if (top_n is None or top_n > 0) else self.top_n
results = self.client.rerank(
query, docs, model, top_n=top_n, max_chunks_per_doc=max_chunks_per_doc
query=query,
documents=docs,
model=model,
top_n=top_n,
max_chunks_per_doc=max_chunks_per_doc,
)
result_dicts = []
for res in results:
for res in results.results:
result_dicts.append(
{"index": res.index, "relevance_score": res.relevance_score}
)

16
libs/partners/cohere/tests/integration_tests/test_rerank.py Normal file
View File

@ -0,0 +1,16 @@
"""Test Cohere reranks."""
from langchain_core.documents import Document
from langchain_cohere import CohereRerank
def test_langchain_cohere_rerank_documents() -> None:
"""Test cohere rerank."""
rerank = CohereRerank()
test_documents = [
Document(page_content="This is a test document."),
Document(page_content="Another test document."),
]
test_query = "Test query"
results = rerank.rerank(test_documents, test_query)
assert len(results) == 2