Skip to content

Add FlashRank reranker to HybridRetriever to improve retrieval quality - #116

Closed
GovindhKishore wants to merge 1 commit into
reactome:mainfrom
GovindhKishore:feature/flashrank-reranking
Closed

GovindhKishore wants to merge 1 commit into
reactome:mainfrom
GovindhKishore:feature/flashrank-reranking

Conversation

@GovindhKishore

Copy link
Copy Markdown

Summary

Adds a reranking layer to HybridRetriever in csv_chroma.py to address the issue of responses becoming increasingly long and noisy as more data sources are integrated into the retrieval pipeline.

Problem

The current pipeline retrieves documents from multiple subdirectories using BM25 + SelfQuery + MultiQuery expansion, resulting in ~90 documents being passed directly to create_stuff_documents_chain.

There is no cross-subdirectory relevance filtering - all retrieved documents are stuffed into the LLM prompt regardless of how relevant they are to the original user query. This causes:

  • Responses becoming longer and noisier as more data is added
  • Low-relevance documents from one subdirectory treated equally to high-relevance documents from another
  • LLM receiving too much context which reduces answer precision

Solution

A new module src/retrievers/reranker.py is introduced using FlashRank (ms-marco-MiniLM-L-12-v2). After weighted_reciprocal_rank merges results across all subdirectories, the reranker scores every retrieved document against the original user query using a cross-encoder model and returns only the top N most relevant documents.

Two functions are provided:

  • rerank() - sync, called by retrieve_documents()
  • arerank() - async, called by aretrieve_documents()

arerank() uses asyncio.to_thread to run the blocking FlashRank inference in a background thread without freezing the async event loop.

Changes

  • src/retrievers/reranker.py - new module containing reranking logic
  • src/retrievers/csv_chroma.py - import reranker, update return statements in both retrieve_documents() and aretrieve_documents()
  • config_default.yml - add reranker configuration block
  • pyproject.toml / poetry.lock - add flashrank dependency

Why FlashRank

  • Runs locally - no API key required
  • CPU only - no GPU needed
  • Lightweight (~4MB model)
  • No changes to downstream pipeline - same list[Document] type
    returned throughout

Impact

Since csv_chroma.py is shared by both Reactome and UniProt retrievers, reranking applies automatically to all current and future database integrations without any additional changes.

Test

# Input: 7 documents (mix of relevant and irrelevant)
# Query: "What does TP53 do in apoptosis?"

# Output after reranking (top 3):
# 1. score=0.9996 | TP53 activates apoptosis through BAX
# 2. score=0.9860 | TP53 and PUMA in intrinsic apoptosis  
# 3. score=0.8930 | p53 regulates cell death signalling

# Correctly dropped:
# RNA polymerase II transcription      (irrelevant)
# Reactome database overview           (irrelevant)
# General cancer pathway summary       (irrelevant)

Note

This contribution was developed with AI assistance (Claude) for understanding the codebase and implementation guidance. All code has been reviewed and understood.

Closes #115

Happy to make any changes based on maintainer feedback.

@GovindhKishore

Copy link
Copy Markdown
Author

Hi @adamjohnwright @GFJHogue ,

Just flagging this PR for your attention when you get a chance. This directly addresses the retrieval noise issue mentioned across several issues, and since it touches csv_chroma.py which is shared by both Reactome and UniProt retrievers, I wanted to make sure the right people are aware of it.

Happy to:

  • Add unit tests if needed
  • Adjust the top_n default value in config_default.yml
  • Discuss alternative reranking models if FlashRank is not preferred

Looking forward to any feedback!

@adamjohnwright

Copy link
Copy Markdown
Contributor

@heliamoh are you able to take a look to see if this resolves the issue(s)?

@adamjohnwright

Copy link
Copy Markdown
Contributor

Still open deliberately, and queued rather than ignored — sorry it has been quiet for six months.

This PR changes what reaches the LLM, and this repository's constitution requires a before-and-after measurement on real questions for exactly that kind of change. "This should be better" is explicitly not a finding here, because retrieval quality has no right answer, only "did this move".

The reason there has been no verdict is that the tool which produces that measurement was not trustworthy. src/evaluation/evaluator.py built its own retriever instead of the shipping pipeline, so it measured a configuration that no longer existed; that is fixed. It then scored the raw question rather than the rephrased one, while production always retrieves on rewritten text; that is fixed too. What remained was that a single failed question discarded the entire run, with nothing written to disk until the very end — which is why nobody wanted to run it.

#211 fixes that last part. Once it lands, this PR gets a real answer: the golden questions, before and after, with the noise floor reported, and the result is a number rather than an opinion.

Two things that will need doing first, so they are not a surprise:

Nothing is needed from you right now. Thank you for the contribution, and for your patience with the delay.

@adamjohnwright

Copy link
Copy Markdown
Contributor

Reviewed against current main. The idea is sound and one part of the implementation is better than it looks, but it cannot land in this shape — and the first reason is decisive before any question of retrieval quality.

It stops the app from starting

Config uses extra="forbid", deliberately: configuration that cannot be honoured must stop the process rather than be silently ignored. Adding a reranker: section to config_default.yml without a matching model gives:

extra_forbidden
Refusing to start. Fix the file, or remove it to run with the defaults in config_default.yml.

.config.schema.yaml and src/util/config_yml/ both need the new section. As written, merging this takes the chatbot down.

The config it adds is never read

reranker:
  enabled: true
  top_n: 5
  model: "ms-marco-MiniLM-L-12-v2"
  cache_dir: "/tmp"
  max_length: 512

reranker.py hardcodes Ranker(model_name="ms-marco-MiniLM-L-12-v2") and takes top_n as a function default. Nothing reads enabled, cache_dir, or max_length. So the knobs look adjustable and aren't — including enabled, which reads like a feature flag and isn't one.

The model loads at import

ranker = Ranker(model_name="ms-marco-MiniLM-L-12-v2")   # module scope

That downloads and loads an ONNX model when the module is first imported, which means it happens during startup, on the network. In the deployed container /tmp is ephemeral, so it re-downloads on every restart; and if the download fails, the import fails and the app does not come up. This box is at 94% disk with ~2.35 GB per chatbot image, so a runtime model download is not free here either.

What is genuinely good

# ranker.rerank() is a blocking CPU operation (neural network inference)
# calling it directly inside async would freeze the entire event loop
results = await asyncio.to_thread(ranker.rerank, request)

That is exactly right, and it is the kind of thing that is usually got wrong. Blocking inference on the event loop would have stalled every concurrent request, and the comment shows it was reasoned about rather than copied.

On whether reranking helps

Unanswered, and it should be answered with a number rather than an argument — this repository requires a before-and-after on real questions for anything that changes what reaches the LLM. The evaluator that produces that is now fixed and hardened (#211), so the measurement is available to whoever picks this up.

Worth noting top_n=5 cuts 40 retrieved documents to 5. That is a large reduction, and plausibly the right one — beta transcripts this week show retrieval returning genuinely irrelevant documents that the model then wrote around. But 8x fewer documents is exactly the kind of change that needs measuring rather than assuming.

Closing because it cannot start the app as written, not because the direction is wrong. Sorry it took six months to get a real review, and thank you for the async care — that part would have been easy to get wrong.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Feature] Improve retrieval quality by adding reranking layer to HybridRetriever

2 participants