feat(retrieval): add pluggable cross-encoder reranker - #93
Open
vcvyg wants to merge 1 commit into
Open
Conversation
Author
There was a problem hiding this comment.
Pull request overview
This PR introduces a pluggable retrieval reranker option by adding a lazily-loaded Sentence Transformers cross-encoder implementation, wiring it into settings-driven configuration, and documenting/benchmarking the new path.
Changes:
- Added
CrossEncoderReranker(batched scoring, lazy dependency import) and exposed it via retrieval module exports. - Added
reranker_mode=cross_encoderconfiguration viaRetrievalSettings+ factory wiring, and injected the reranker intoContextSeekretrieval. - Added unit tests, documentation updates, and a small benchmark script for recall/latency comparison.
Reviewed changes
Copilot reviewed 16 out of 17 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| uv.lock | Adds sentence-transformers and updates lock resolution to support the new optional extra. |
| pyproject.toml | Introduces rerank extra and includes sentence-transformers in all. |
| src/contextseek/retrieval/components.py | Implements CrossEncoderReranker with lazy model loading and batched prediction. |
| src/contextseek/retrieval/init.py | Re-exports CrossEncoderReranker from retrieval package. |
| src/contextseek/config/settings.py | Adds cross-encoder configuration fields to RetrievalSettings. |
| src/contextseek/config/factory.py | Adds build_reranker() to construct non-LLM rerankers based on settings. |
| src/contextseek/config/init.py | Exports build_reranker for external configuration assembly. |
| src/contextseek/config/strategies.py | Updates comment/documentation for reranker_mode in RetrievalStrategy. |
| src/contextseek/client/contextseek.py | Injects configured reranker into retrieval; preserves existing LLM rerank behavior. |
| tests/unit_tests/test_settings.py | Adds settings/factory coverage for cross-encoder reranker selection and validation. |
| tests/unit_tests/retrieval/test_reranker_features.py | Adds stub-backed unit tests for cross-encoder reranking behavior and fallback. |
| scripts/benchmark_rerankers.py | Adds reproducible benchmark script for heuristic vs cross-encoder recall/latency. |
| docs/en/reference/settings.md | Documents new RETRIEVAL_CROSS_ENCODER_* settings and mode. |
| docs/zh/reference/settings.md | Chinese settings reference updates for cross-encoder mode/settings. |
| docs/en/concepts/retrieval-model.md | Documents cross-encoder reranker usage and configuration. |
| docs/en/concepts/reranker-benchmark.md | Adds benchmark note and reference results table. |
| .env.example | Adds example env vars for cross-encoder reranker configuration. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+807
to
+814
| try: | ||
| scores = list(model.predict(pairs)) | ||
| if len(scores) != len(to_score): | ||
| raise ValueError("cross-encoder returned an unexpected score count") | ||
| for item, score in zip(to_score, scores): | ||
| item["_score"] = round(float(score), 6) | ||
| except Exception: # noqa: BLE001 | ||
| return pre_ranked |
Comment on lines
+68
to
71
| # Rerank mode: "heuristic" (default), "llm", or "cross_encoder" | ||
| reranker_mode: str = "heuristic" | ||
| # Limit number of candidates scored by LLM in reranking | ||
| llm_rerank_top_n: int = 20 |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
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.
Summary
CrossEncoderRerankerthat implements the existing reranker protocol and lazily loads the optional Sentence Transformers dependencyheuristic,llm, andcross_encoderselectable through retrieval settings while preserving the existing LLM reranking pathValidation
uv run ruff check src/contextseek tests/unit_tests scripts/benchmark_rerankers.pyuv run pytest tests/ -q --ignore=tests/unit_tests/test_powermem_plug.py(451 passed, 8 skipped)uv run --extra rerank python scripts/benchmark_rerankers.py --device cpu --rounds 5Closes #49