feat: Add Cohere reranker config, chunking, and tests #2411
+620
−20
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Description
This pull request introduces support for automatic document chunking in the reranking pipeline, which is especially useful for models with strict token limits (such as ColBERT or Cohere rerank models). The changes allow documents that exceed the token limit to be split into manageable chunks, reranked individually, and then aggregated back to provide a final relevance score per original document. This enhancement is configurable via environment variables and is integrated into both the main API and example scripts.
Related Issues
N/A
Changes Made
chunk_documents_for_rerankandaggregate_chunk_scoresutilities tolightrag/rerank.pyto handle splitting long documents into token-limited chunks and aggregating their rerank scores back to the original documents. This includes both a tokenizer-based and a character-based fallback approach.generic_rerank_apiandcohere_rerankfunctions to accept and process chunking parametersAPI and Server Integration
lightrag/api/lightrag_server.pyto read chunking configuration from environment variables and pass these options to the rerank function when using the Cohere binding.Configuration and Examples
RERANK_ENABLE_CHUNKING,RERANK_MAX_TOKENS_PER_DOC) toenv.examplefor configuring chunking behavior.examples/rerank_example.pyto document and utilize the new chunking configuration options, including reasonable defaults and usage notes.Checklist
Additional Notes
[Add any additional notes or context for the reviewer(s).]