Files

2.7 KiB

rerank

Companion code for Reranking RAG Results in Java: Cross-Encoder Reranker with ONNX Runtime Before the LLM.

A bi-encoder (all-MiniLM-L6-v2) retrieves candidates, a cross-encoder (ms-marco-MiniLM-L-6-v2) rescores them, and a Spring AI DocumentPostProcessor keeps the top 5 for RetrievalAugmentationAdvisor. Both models run in-process through ONNX Runtime; the tokenizers are DJL's Hugging Face tokenizers. No database, no API key, no GPU.

Versions
Spring Boot 4.1.2
Spring AI 2.0.2
ONNX Runtime 1.31.0 (the spring-ai-transformers starter would bring 1.21.1)
DJL tokenizers 0.39.0 (the starter would bring 0.36.0)
JDK 25

Layout

File What it is
BiEncoder ONNX Runtime + tokenizer + mean pooling: text to a 384-dim vector
BiEncoderRetriever in-memory cosine search, a Spring AI DocumentRetriever
CrossEncoder [CLS] query [SEP] passage [SEP] through the model, one logit out
RerankingPostProcessor the DocumentPostProcessor: score, sort, keep top N, record both ranks in metadata
Models where the model files come from (Hugging Face, cached in ~/.cache/rerank)
KnowledgeBase the same 94 passages and 60 labelled questions as the hybrid-search module

Run

scripts/run-all.sh. The first run downloads about 200 MB of model files. Every file in output/ is written by a test that asserts the same facts. Files 06 and 07 contain wall-clock latency and depend on your machine; the others were identical across two consecutive runs.

Output Shows
01-pair-encoding.txt the tokens, token_type_ids and attention mask of a (query, passage) pair
02-one-question-by-hand.txt one question's top 8 with cosine and cross-encoder logit side by side
03-quality-by-pool-size.txt hit@1/3/5 and MRR for pools of 5, 10, 20, 50 and 94, and first-stage recall
04-quality-by-question-kind.txt the same, split by exact code, SKU, paraphrase and mixed questions
05-helped-and-hurt.txt every question whose gold-passage rank changed
06-latency-by-candidates-and-batch.txt rerank time by candidate count, batch size and thread count
07-latency-against-quality.txt fp32 and int8, quality and latency in one table
08-post-processor-in-the-advisor.txt retrieve 50, keep 5, and the prompt the model receives
09-post-processor-edge-cases.txt empty input, fewer documents than N
10-14 -trap-*.txt missing input, swapped pair order, score scale, truncation, pool too small

The corpus is synthetic and small. It shows how the methods behave, not how they scale.