Add rerank module: bi-encoder retrieval plus an ONNX Runtime cross-encoder reranker behind a Spring AI DocumentPostProcessor; hit@k and MRR by pool size and question kind, latency by candidates, batch size, threads and int8, and five reproduced traps. Local MiniLM models, no database or API key.
Co-Authored-By: Claude Sonnet 5.5 <[email protected]> Claude-Session: https://claude.ai/code/session_01Cfo2K4u4rF18QsPXPFnytf
This commit is contained in:
Claude
@@ -0,0 +1 @@
|
||||
target/
|
||||
@@ -0,0 +1,43 @@
|
||||
# rerank
|
||||
|
||||
Companion code for [Reranking RAG Results in Java: Cross-Encoder Reranker with ONNX Runtime Before the LLM](https://ankurm.com/reranking-rag-results-java-cross-encoder-onnx-runtime-spring-ai/).
|
||||
|
||||
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 250 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.
|
||||
@@ -0,0 +1,6 @@
|
||||
# A (query, passage) pair as the cross-encoder receives it
|
||||
|
||||
model inputs: [input_ids, attention_mask, token_type_ids]
|
||||
tokens : [[CLS], what, is, e, -, 412, ##6, ?, [SEP], error, e, -, 412, ##6, -, ref, ##und, window, closed, ., ref, ##unds, can, be, issued, for, 120, days, ., [SEP]]
|
||||
token_type_ids: [0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]
|
||||
attention_mask: [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]
|
||||
@@ -0,0 +1,16 @@
|
||||
# One question, the bi-encoder's top 8, rescored
|
||||
|
||||
question: A buyer told their bank they never got the item. What now?
|
||||
gold passage: G-04
|
||||
|
||||
retrieved id cosine cross-encoder logit
|
||||
1 G-06 0.446 -7.40
|
||||
2 G-04 0.440 -8.06
|
||||
3 E-4117 0.418 -11.08
|
||||
4 E-4101 0.418 -11.30
|
||||
5 G-03 0.405 -10.33
|
||||
6 E-4140 0.403 -9.91
|
||||
7 G-08 0.399 -9.62
|
||||
8 G-24 0.391 -9.61
|
||||
|
||||
bi-encoder rank of G-04: 2; cross-encoder winner: G-06
|
||||
@@ -0,0 +1,16 @@
|
||||
# Rerank the bi-encoder's top N (60 questions, 94 passages, fp32 cross-encoder)
|
||||
|
||||
ranking hit@1 hit@3 hit@5 MRR@10
|
||||
bi-encoder alone 41/60 51/60 55/60 0.782
|
||||
bi-encoder top 5 + cross-encoder 48/60 55/60 55/60 0.863
|
||||
bi-encoder top 10 + cross-encoder 49/60 54/60 56/60 0.868
|
||||
bi-encoder top 20 + cross-encoder 49/60 55/60 56/60 0.870
|
||||
bi-encoder top 50 + cross-encoder 49/60 56/60 56/60 0.869
|
||||
bi-encoder top 94 + cross-encoder 49/60 56/60 56/60 0.869
|
||||
|
||||
recall of the first stage (is the gold passage inside the pool at all?):
|
||||
top 5 : 55/60
|
||||
top 10 : 58/60
|
||||
top 20 : 59/60
|
||||
top 50 : 60/60
|
||||
top 94 : 60/60
|
||||
@@ -0,0 +1,11 @@
|
||||
# Bi-encoder alone against bi-encoder top 50 + cross-encoder, by kind of question
|
||||
|
||||
kind ranking hit@1 hit@3 hit@5 MRR@10
|
||||
exact error code bi-encoder 2/12 5/12 8/12 0.374
|
||||
reranked 8/12 12/12 12/12 0.792
|
||||
exact SKU or setting bi-encoder 12/12 12/12 12/12 1.000
|
||||
reranked 12/12 12/12 12/12 1.000
|
||||
paraphrase bi-encoder 15/24 22/24 23/24 0.767
|
||||
reranked 17/24 20/24 20/24 0.778
|
||||
code plus words bi-encoder 12/12 12/12 12/12 1.000
|
||||
reranked 12/12 12/12 12/12 1.000
|
||||
@@ -0,0 +1,26 @@
|
||||
# Gold passage rank before and after reranking the top 50 (only questions whose rank changed)
|
||||
|
||||
gold question bi-encoder reranked
|
||||
E-4106 What does error E-4106 mean? 28 3 better
|
||||
E-4116 Why am I getting E-4116? 1 2 worse
|
||||
E-4118 We see E-4118 in the logs, what~ 8 3 better
|
||||
E-4120 Explain E-4120 2 1 better
|
||||
E-4122 Support told me to look up E-41~ 4 3 better
|
||||
E-4126 What is E-4126? 2 1 better
|
||||
E-4129 How do I fix E-4129? 12 1 better
|
||||
E-4131 Our service started returning E~ 4 1 better
|
||||
E-4133 Meaning of E-4133 5 1 better
|
||||
E-4135 Is E-4135 our fault? 2 1 better
|
||||
E-4138 What causes E-4138 on the shop ~ 6 1 better
|
||||
G-01 What happens when a customer's ~ 6 11 worse
|
||||
G-05 How do I experiment safely befo~ 1 14 worse
|
||||
G-07 Accept money from people who do~ 4 1 better
|
||||
G-10 Stop dodgy purchases automatica~ 3 6 worse
|
||||
G-12 What if our website is down whe~ 2 1 better
|
||||
G-16 A client wants to stop being bi~ 3 1 better
|
||||
G-17 Give our bookkeeper everything ~ 1 25 worse
|
||||
G-20 How do we avoid ever handling r~ 2 1 better
|
||||
G-24 One item in the basket is sold ~ 3 1 better
|
||||
G-28 Which colleague altered the set~ 1 2 worse
|
||||
|
||||
better: 15, worse: 6, unchanged: 39
|
||||
@@ -0,0 +1,10 @@
|
||||
# Time to rerank N candidates (median and slowest of 12 calls, after 3 warm-up calls)
|
||||
|
||||
machine: 2 vCPU, JDK 25.0.5
|
||||
fp32 cross-encoder, max length 256 tokens; the first column of each cell is the median in ms, the second the slowest
|
||||
|
||||
candidates batch 1 batch 8 batch 32 batch 32, 1 thread
|
||||
10 125 / 188 121 / 142 122 / 138 148 / 208
|
||||
20 205 / 226 195 / 220 220 / 257 322 / 411
|
||||
50 546 / 661 512 / 578 489 / 549 800 / 913
|
||||
94 897 / 1062 838 / 966 962 / 1075 1635 / 1928
|
||||
@@ -0,0 +1,17 @@
|
||||
# Quality and added latency for each setting (60 questions; latency is the median of 12 calls)
|
||||
|
||||
machine: 2 vCPU, JDK 25.0.5
|
||||
first stage (embed the question, compare with 94 vectors): median 3.8 ms
|
||||
|
||||
setting hit@1 hit@3 hit@5 MRR@10 rerank ms (median)
|
||||
no reranker 41/60 51/60 55/60 0.782 0
|
||||
fp32, top 5 48/60 55/60 55/60 0.863 70
|
||||
int8, top 5 48/60 55/60 55/60 0.860 68
|
||||
fp32, top 10 49/60 54/60 56/60 0.868 116
|
||||
int8, top 10 49/60 54/60 56/60 0.868 106
|
||||
fp32, top 20 49/60 55/60 56/60 0.870 199
|
||||
int8, top 20 49/60 55/60 56/60 0.870 197
|
||||
fp32, top 50 49/60 56/60 56/60 0.869 461
|
||||
int8, top 50 49/60 56/60 56/60 0.869 503
|
||||
|
||||
int8 and fp32 put the same passage first for 58 of 60 questions (top 50)
|
||||
@@ -0,0 +1,34 @@
|
||||
# RerankingPostProcessor inside RetrievalAugmentationAdvisor (retrieve 50, keep 5)
|
||||
|
||||
question: How do I fix E-4129?
|
||||
retriever returned 50 documents; the post-processor kept 5
|
||||
|
||||
kept id retrieval rank rerank score text
|
||||
1 E-4129 12 3.52 Error E-4129 - API key lacks permission. The key
|
||||
2 E-4128 2 2.21 Error E-4128 - Tax identifier invalid. The tax r
|
||||
3 E-4122 1 1.48 Error E-4122 - Payment method not attached. The
|
||||
4 E-4123 14 1.09 Error E-4123 - Duplicate customer email. A custo
|
||||
5 E-4126 5 0.94 Error E-4126 - Coupon expired. The coupon passed
|
||||
|
||||
--- the prompt the model received ---
|
||||
Context information is below.
|
||||
|
||||
---------------------
|
||||
Error E-4129 - API key lacks permission. The key in use does not carry the scope needed for this endpoint. Create a restricted key that includes it.
|
||||
Error E-4128 - Tax identifier invalid. The tax registration number does not match the format for its country. Check the prefix and the number of digits.
|
||||
Error E-4122 - Payment method not attached. The payment method is not attached to this customer. Attach it first, or pass it directly on a one-off charge.
|
||||
Error E-4123 - Duplicate customer email. A customer with this email address already exists and the account is configured to enforce unique emails. Reuse the existing customer.
|
||||
Error E-4126 - Coupon expired. The coupon passed its redemption end date. Create a new coupon or extend the validity in the dashboard.
|
||||
---------------------
|
||||
|
||||
Given the context information and no prior knowledge, answer the query.
|
||||
|
||||
Follow these rules:
|
||||
|
||||
1. If the answer is not in the context, just say that you don't know.
|
||||
2. Avoid statements like "Based on the context..." or "The provided information...".
|
||||
|
||||
Query: How do I fix E-4129?
|
||||
|
||||
Answer:
|
||||
|
||||
@@ -0,0 +1,4 @@
|
||||
# RerankingPostProcessor edge cases
|
||||
|
||||
no documents in -> 0 out
|
||||
3 documents in, topN = 5 -> 3 out (nothing invented, nothing dropped)
|
||||
@@ -0,0 +1,6 @@
|
||||
# Leaving out token_type_ids
|
||||
|
||||
model inputs : [input_ids, attention_mask, token_type_ids]
|
||||
we sent : [input_ids, attention_mask]
|
||||
result : Error code - ORT_FAIL - message: Non-zero status code returned while running Gather node. Name:'/bert/embeddings/token_type_embeddings/Gather' Status Message: /onnxruntime_src/include/onnxruntime/core/framework/op_kernel_context.h:45 const T* onnxruntime::OpKernelContext::Input(int) const [with T = onnxruntime::Tensor] Missing Input: token_type_ids
|
||||
|
||||
@@ -0,0 +1,5 @@
|
||||
# (passage, query) instead of (query, passage); top 20 reranked
|
||||
|
||||
(query, passage) hit@1 49/60 MRR 0.870
|
||||
(passage, query) hit@1 37/60 MRR 0.708
|
||||
no error either way: both orders return a number for every pair
|
||||
@@ -0,0 +1,8 @@
|
||||
# Raw cross-encoder scores for the gold passage and for the best wrong passage
|
||||
|
||||
gold passage score : min -11.5, median 2.1, max 9.0
|
||||
best wrong passage score : min -11.4, median -1.4, max 7.8
|
||||
questions whose best-scoring passage is below 0: 26 of 60
|
||||
a fixed cut-off of 0 would drop the gold passage for 26 of 60 questions
|
||||
a fixed cut-off of 2 would drop the gold passage for 30 of 60 questions
|
||||
a fixed cut-off of 4 would drop the gold passage for 38 of 60 questions
|
||||
@@ -0,0 +1,7 @@
|
||||
# The answer sits at the end of a long passage
|
||||
|
||||
question + passage = 224 tokens
|
||||
score with max length 512 : -1.50
|
||||
score with max length 64 : -11.03 (the answer was cut off before the model saw it)
|
||||
score of the answer alone : 2.46
|
||||
passages in this corpus longer than 256 tokens: 0 of 94
|
||||
@@ -0,0 +1,11 @@
|
||||
# Questions whose gold passage is not in the bi-encoder's top 5
|
||||
|
||||
gold question bi-encoder reranked top 50
|
||||
E-4106 What does error E-4106 mean? 28 3
|
||||
E-4118 We see E-4118 in the logs, what should we do? 8 3
|
||||
E-4129 How do I fix E-4129? 12 1
|
||||
E-4138 What causes E-4138 on the shop floor? 6 1
|
||||
G-01 What happens when a customer's monthly payment does not go ~ 6 11
|
||||
|
||||
5 questions have no answer passage in the top 5; a reranker over the top 5 cannot help them,
|
||||
a reranker over the top 50 brings 4 of them into the top 5
|
||||
@@ -0,0 +1,65 @@
|
||||
<?xml version="1.0" encoding="UTF-8"?>
|
||||
<project xmlns="http://maven.apache.org/POM/4.0.0"
|
||||
xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
|
||||
xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 https://maven.apache.org/xsd/maven-4.0.0.xsd">
|
||||
<modelVersion>4.0.0</modelVersion>
|
||||
|
||||
<parent>
|
||||
<groupId>org.springframework.boot</groupId>
|
||||
<artifactId>spring-boot-starter-parent</artifactId>
|
||||
<version>4.1.2</version>
|
||||
<relativePath/>
|
||||
</parent>
|
||||
|
||||
<groupId>com.ankurm</groupId>
|
||||
<artifactId>rerank</artifactId>
|
||||
<version>1.0.0</version>
|
||||
<name>rerank</name>
|
||||
<description>Cross-encoder reranking for RAG in Java: a bi-encoder retrieves 50 candidates, an ONNX Runtime cross-encoder rescores them, a Spring AI DocumentPostProcessor keeps the top 5. Quality and latency measured on 60 labelled queries.</description>
|
||||
|
||||
<properties>
|
||||
<java.version>25</java.version>
|
||||
<spring-ai.version>2.0.2</spring-ai.version>
|
||||
<!-- Latest releases on Maven Central on 2026-10-24. -->
|
||||
<onnxruntime.version>1.31.0</onnxruntime.version>
|
||||
<djl.version>0.39.0</djl.version>
|
||||
</properties>
|
||||
|
||||
<dependencyManagement>
|
||||
<dependencies>
|
||||
<dependency>
|
||||
<groupId>org.springframework.ai</groupId>
|
||||
<artifactId>spring-ai-bom</artifactId>
|
||||
<version>${spring-ai.version}</version>
|
||||
<type>pom</type>
|
||||
<scope>import</scope>
|
||||
</dependency>
|
||||
</dependencies>
|
||||
</dependencyManagement>
|
||||
|
||||
<dependencies>
|
||||
<!-- DocumentPostProcessor, DocumentRetriever, Query and RetrievalAugmentationAdvisor. -->
|
||||
<dependency><groupId>org.springframework.ai</groupId><artifactId>spring-ai-rag</artifactId></dependency>
|
||||
<!-- Deliberately NOT spring-ai-transformers: its TransformersEmbeddingModel needs a DJL engine that downloads a PyTorch native library at first use. ONNX Runtime and the DJL tokenizers are enough for both stages. -->
|
||||
<dependency><groupId>com.microsoft.onnxruntime</groupId><artifactId>onnxruntime</artifactId><version>${onnxruntime.version}</version></dependency>
|
||||
<dependency><groupId>ai.djl</groupId><artifactId>api</artifactId><version>${djl.version}</version></dependency>
|
||||
<dependency><groupId>ai.djl.huggingface</groupId><artifactId>tokenizers</artifactId><version>${djl.version}</version></dependency>
|
||||
<dependency>
|
||||
<groupId>org.springframework.boot</groupId>
|
||||
<artifactId>spring-boot-starter-test</artifactId>
|
||||
<scope>test</scope>
|
||||
</dependency>
|
||||
</dependencies>
|
||||
|
||||
<build>
|
||||
<plugins>
|
||||
<plugin>
|
||||
<groupId>org.apache.maven.plugins</groupId>
|
||||
<artifactId>maven-surefire-plugin</artifactId>
|
||||
<configuration>
|
||||
<argLine>-Duser.timezone=UTC -Dstdout.encoding=UTF-8 -Dfile.encoding=UTF-8 -Xmx3g</argLine>
|
||||
</configuration>
|
||||
</plugin>
|
||||
</plugins>
|
||||
</build>
|
||||
</project>
|
||||
Executable
+9
@@ -0,0 +1,9 @@
|
||||
#!/usr/bin/env bash
|
||||
# Regenerates every file in output/ by running the test suite.
|
||||
# Needs: JDK 25, Maven, and internet on the first run (about 250 MB of models from Hugging Face, cached in ~/.cache/rerank).
|
||||
# No database, no API key. Latency files (06, 07) depend on your machine; every other file is deterministic.
|
||||
set -euo pipefail
|
||||
cd "$(dirname "$0")/.."
|
||||
rm -rf output
|
||||
mvn -B -q test
|
||||
ls output
|
||||
@@ -0,0 +1,109 @@
|
||||
package com.ankurm.rerank;
|
||||
|
||||
import java.nio.file.Path;
|
||||
import java.util.ArrayList;
|
||||
import java.util.HashMap;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
|
||||
import ai.djl.huggingface.tokenizers.Encoding;
|
||||
import ai.djl.huggingface.tokenizers.HuggingFaceTokenizer;
|
||||
import ai.onnxruntime.OnnxTensor;
|
||||
import ai.onnxruntime.OrtEnvironment;
|
||||
import ai.onnxruntime.OrtException;
|
||||
import ai.onnxruntime.OrtSession;
|
||||
|
||||
/**
|
||||
* The retrieval-stage model: text in, one 384-number vector out, computed for each text on its own. The
|
||||
* ONNX file returns one vector per token; sentence-transformers' recipe for all-MiniLM-L6-v2 is to average
|
||||
* them (ignoring padding) and scale the result to length 1, which is what {@link #embedAll} does.
|
||||
*/
|
||||
public final class BiEncoder implements AutoCloseable {
|
||||
|
||||
private final OrtEnvironment env = OrtEnvironment.getEnvironment();
|
||||
private final OrtSession session;
|
||||
private final HuggingFaceTokenizer tokenizer;
|
||||
|
||||
public BiEncoder(Path model, Path tokenizerJson, int maxLength) {
|
||||
try {
|
||||
this.session = env.createSession(model.toString(), new OrtSession.SessionOptions());
|
||||
this.tokenizer = HuggingFaceTokenizer.builder().optTokenizerPath(tokenizerJson)
|
||||
.optTruncation(true).optPadding(true).optMaxLength(maxLength).build();
|
||||
} catch (Exception e) {
|
||||
throw new IllegalStateException(e);
|
||||
}
|
||||
}
|
||||
|
||||
public float[] embed(String text) {
|
||||
return embedAll(List.of(text)).get(0);
|
||||
}
|
||||
|
||||
public List<float[]> embedAll(List<String> texts) {
|
||||
List<float[]> out = new ArrayList<>();
|
||||
for (int from = 0; from < texts.size(); from += 32) {
|
||||
out.addAll(run(texts.subList(from, Math.min(from + 32, texts.size())).toArray(String[]::new)));
|
||||
}
|
||||
return out;
|
||||
}
|
||||
|
||||
private List<float[]> run(String[] texts) {
|
||||
Encoding[] enc = tokenizer.batchEncode(texts);
|
||||
int n = enc.length;
|
||||
long[][] ids = new long[n][];
|
||||
long[][] mask = new long[n][];
|
||||
long[][] types = new long[n][];
|
||||
for (int i = 0; i < n; i++) {
|
||||
ids[i] = enc[i].getIds();
|
||||
mask[i] = enc[i].getAttentionMask();
|
||||
types[i] = enc[i].getTypeIds();
|
||||
}
|
||||
Map<String, OnnxTensor> in = new HashMap<>();
|
||||
try {
|
||||
in.put("input_ids", OnnxTensor.createTensor(env, ids));
|
||||
in.put("attention_mask", OnnxTensor.createTensor(env, mask));
|
||||
in.put("token_type_ids", OnnxTensor.createTensor(env, types));
|
||||
try (OrtSession.Result r = session.run(in)) {
|
||||
float[][][] hidden = (float[][][]) r.get(0).getValue();
|
||||
List<float[]> out = new ArrayList<>();
|
||||
for (int i = 0; i < n; i++) {
|
||||
int dim = hidden[i][0].length;
|
||||
float[] v = new float[dim];
|
||||
double count = 0;
|
||||
for (int t = 0; t < hidden[i].length; t++) {
|
||||
if (mask[i][t] == 1) {
|
||||
count++;
|
||||
for (int d = 0; d < dim; d++) {
|
||||
v[d] += hidden[i][t][d];
|
||||
}
|
||||
}
|
||||
}
|
||||
double norm = 0;
|
||||
for (int d = 0; d < dim; d++) {
|
||||
v[d] /= count;
|
||||
norm += v[d] * v[d];
|
||||
}
|
||||
norm = Math.sqrt(norm);
|
||||
for (int d = 0; d < dim; d++) {
|
||||
v[d] /= (float) norm;
|
||||
}
|
||||
out.add(v);
|
||||
}
|
||||
return out;
|
||||
}
|
||||
} catch (OrtException e) {
|
||||
throw new IllegalStateException(e.getMessage(), e);
|
||||
} finally {
|
||||
in.values().forEach(OnnxTensor::close);
|
||||
}
|
||||
}
|
||||
|
||||
@Override
|
||||
public void close() {
|
||||
try {
|
||||
session.close();
|
||||
} catch (OrtException e) {
|
||||
throw new IllegalStateException(e);
|
||||
}
|
||||
tokenizer.close();
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,69 @@
|
||||
package com.ankurm.rerank;
|
||||
|
||||
import java.util.Comparator;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
import java.util.stream.IntStream;
|
||||
|
||||
import org.springframework.ai.document.Document;
|
||||
import org.springframework.ai.rag.Query;
|
||||
import org.springframework.ai.rag.retrieval.search.DocumentRetriever;
|
||||
|
||||
/**
|
||||
* The first stage: embed every passage once, embed the question, rank by cosine similarity. Held in memory
|
||||
* so the module needs no database; a vector store does the same job with an index. The retrieval rank and
|
||||
* score travel with each Document so the reranker's effect can be read off afterwards.
|
||||
*/
|
||||
public class BiEncoderRetriever implements DocumentRetriever {
|
||||
|
||||
private final BiEncoder model;
|
||||
private final List<KnowledgeBase.Passage> passages;
|
||||
private final float[][] vectors;
|
||||
private final int topK;
|
||||
|
||||
public BiEncoderRetriever(BiEncoder model, List<KnowledgeBase.Passage> passages, int topK) {
|
||||
this.model = model;
|
||||
this.passages = passages;
|
||||
this.topK = topK;
|
||||
this.vectors = model.embedAll(passages.stream().map(KnowledgeBase.Passage::text).toList()).toArray(new float[0][]);
|
||||
}
|
||||
|
||||
public BiEncoderRetriever withTopK(int k) {
|
||||
return new BiEncoderRetriever(this, k);
|
||||
}
|
||||
|
||||
private BiEncoderRetriever(BiEncoderRetriever o, int k) {
|
||||
this.model = o.model;
|
||||
this.passages = o.passages;
|
||||
this.vectors = o.vectors;
|
||||
this.topK = k;
|
||||
}
|
||||
|
||||
@Override
|
||||
public List<Document> retrieve(Query query) {
|
||||
float[] q = model.embed(query.text());
|
||||
double[] sim = new double[vectors.length];
|
||||
for (int i = 0; i < vectors.length; i++) {
|
||||
sim[i] = cosine(q, vectors[i]);
|
||||
}
|
||||
List<Integer> order = IntStream.range(0, vectors.length).boxed()
|
||||
.sorted(Comparator.comparingDouble((Integer i) -> -sim[i])).limit(topK).toList();
|
||||
java.util.ArrayList<Document> out = new java.util.ArrayList<>();
|
||||
for (int rank = 0; rank < order.size(); rank++) {
|
||||
KnowledgeBase.Passage p = passages.get(order.get(rank));
|
||||
out.add(Document.builder().id(p.pid()).text(p.text()).score(sim[order.get(rank)])
|
||||
.metadata(Map.of("pid", p.pid(), "retrieval_rank", rank + 1)).build());
|
||||
}
|
||||
return out;
|
||||
}
|
||||
|
||||
static double cosine(float[] a, float[] b) {
|
||||
double dot = 0, na = 0, nb = 0;
|
||||
for (int i = 0; i < a.length; i++) {
|
||||
dot += a[i] * b[i];
|
||||
na += a[i] * a[i];
|
||||
nb += b[i] * b[i];
|
||||
}
|
||||
return dot / Math.sqrt(na * nb);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,124 @@
|
||||
package com.ankurm.rerank;
|
||||
|
||||
import java.nio.file.Path;
|
||||
import java.util.ArrayList;
|
||||
import java.util.HashMap;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
|
||||
import ai.djl.huggingface.tokenizers.Encoding;
|
||||
import ai.djl.huggingface.tokenizers.HuggingFaceTokenizer;
|
||||
import ai.djl.util.PairList;
|
||||
import ai.onnxruntime.OnnxTensor;
|
||||
import ai.onnxruntime.OrtEnvironment;
|
||||
import ai.onnxruntime.OrtException;
|
||||
import ai.onnxruntime.OrtSession;
|
||||
|
||||
/**
|
||||
* A cross-encoder: the query and one passage go into the model together, as a single sequence
|
||||
* {@code [CLS] query [SEP] passage [SEP]}, and one number comes out. Unlike a bi-encoder, nothing about a
|
||||
* passage can be computed ahead of time, which is why it is only affordable on a short candidate list.
|
||||
*
|
||||
* <p>The number is a raw logit. Higher is more relevant. It is not a probability and it is only meaningful
|
||||
* for ordering passages <em>for the same query</em>.
|
||||
*/
|
||||
public final class CrossEncoder implements AutoCloseable {
|
||||
|
||||
private final OrtEnvironment env = OrtEnvironment.getEnvironment();
|
||||
private final OrtSession session;
|
||||
private final HuggingFaceTokenizer tokenizer;
|
||||
private final int batchSize;
|
||||
private final boolean sendTypeIds;
|
||||
|
||||
public CrossEncoder(Path model, Path tokenizerJson, int maxLength, int batchSize, int intraOpThreads) {
|
||||
this(model, tokenizerJson, maxLength, batchSize, intraOpThreads, true);
|
||||
}
|
||||
|
||||
/** {@code sendTypeIds=false} exists only so a test can show what a missing input looks like. */
|
||||
public CrossEncoder(Path model, Path tokenizerJson, int maxLength, int batchSize, int intraOpThreads, boolean sendTypeIds) {
|
||||
try {
|
||||
OrtSession.SessionOptions opts = new OrtSession.SessionOptions();
|
||||
if (intraOpThreads > 0) {
|
||||
opts.setIntraOpNumThreads(intraOpThreads);
|
||||
}
|
||||
this.session = env.createSession(model.toString(), opts);
|
||||
this.tokenizer = HuggingFaceTokenizer.builder().optTokenizerPath(tokenizerJson)
|
||||
.optTruncation(true).optPadding(true).optMaxLength(maxLength).build();
|
||||
} catch (Exception e) {
|
||||
throw new IllegalStateException(e);
|
||||
}
|
||||
this.batchSize = batchSize;
|
||||
this.sendTypeIds = sendTypeIds;
|
||||
}
|
||||
|
||||
public List<String> inputNames() {
|
||||
return new ArrayList<>(session.getInputNames());
|
||||
}
|
||||
|
||||
/** The tokenised (query, passage) pair, exactly as it is fed to the model. */
|
||||
public Encoding encode(String query, String passage) {
|
||||
return tokenizer.encode(query, passage);
|
||||
}
|
||||
|
||||
/** One score per passage, in the order given. */
|
||||
public float[] score(String query, List<String> passages) {
|
||||
float[] out = new float[passages.size()];
|
||||
for (int from = 0; from < passages.size(); from += batchSize) {
|
||||
int to = Math.min(from + batchSize, passages.size());
|
||||
PairList<String, String> pairs = new PairList<>();
|
||||
for (String p : passages.subList(from, to)) {
|
||||
pairs.add(query, p);
|
||||
}
|
||||
Encoding[] enc = tokenizer.batchEncode(pairs);
|
||||
float[] part = run(enc);
|
||||
System.arraycopy(part, 0, out, from, part.length);
|
||||
}
|
||||
return out;
|
||||
}
|
||||
|
||||
private float[] run(Encoding[] enc) {
|
||||
int n = enc.length;
|
||||
int len = enc[0].getIds().length;
|
||||
long[][] ids = new long[n][];
|
||||
long[][] mask = new long[n][];
|
||||
long[][] types = new long[n][];
|
||||
for (int i = 0; i < n; i++) {
|
||||
ids[i] = enc[i].getIds();
|
||||
mask[i] = enc[i].getAttentionMask();
|
||||
types[i] = enc[i].getTypeIds();
|
||||
if (ids[i].length != len) {
|
||||
throw new IllegalStateException("batch is not padded to one length");
|
||||
}
|
||||
}
|
||||
Map<String, OnnxTensor> inputs = new HashMap<>();
|
||||
try {
|
||||
inputs.put("input_ids", OnnxTensor.createTensor(env, ids));
|
||||
inputs.put("attention_mask", OnnxTensor.createTensor(env, mask));
|
||||
if (sendTypeIds) {
|
||||
inputs.put("token_type_ids", OnnxTensor.createTensor(env, types));
|
||||
}
|
||||
try (OrtSession.Result r = session.run(inputs)) {
|
||||
float[][] logits = (float[][]) r.get(0).getValue();
|
||||
float[] out = new float[n];
|
||||
for (int i = 0; i < n; i++) {
|
||||
out[i] = logits[i][0];
|
||||
}
|
||||
return out;
|
||||
}
|
||||
} catch (OrtException e) {
|
||||
throw new IllegalStateException(e.getMessage(), e);
|
||||
} finally {
|
||||
inputs.values().forEach(OnnxTensor::close);
|
||||
}
|
||||
}
|
||||
|
||||
@Override
|
||||
public void close() {
|
||||
try {
|
||||
session.close();
|
||||
} catch (OrtException e) {
|
||||
throw new IllegalStateException(e);
|
||||
}
|
||||
tokenizer.close();
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,222 @@
|
||||
package com.ankurm.rerank;
|
||||
|
||||
import java.util.ArrayList;
|
||||
import java.util.List;
|
||||
|
||||
/**
|
||||
* A made-up support knowledge base for a payments API called Orbit: 94 passages and 60 labelled queries.
|
||||
* Every query has exactly one passage that answers it, named by {@code pid}.
|
||||
*
|
||||
* <p>The passages come in three kinds, because the three kinds are where the two search methods differ:
|
||||
* 40 error codes whose texts all look alike and differ in the identifier and one fact, 24 hardware and
|
||||
* configuration identifiers (SKUs and setting names), and 30 how-to guides written in plain language.
|
||||
* The corpus is synthetic and small. It shows how the methods behave, not how any of them scales.
|
||||
*/
|
||||
public final class KnowledgeBase {
|
||||
|
||||
public record Passage(String pid, String kind, String text) {
|
||||
}
|
||||
|
||||
public record Query(String text, String pid, String category) {
|
||||
}
|
||||
|
||||
public static final String EXACT_CODE = "exact error code";
|
||||
public static final String EXACT_ID = "exact SKU or setting";
|
||||
public static final String PARAPHRASE = "paraphrase";
|
||||
public static final String MIXED = "code plus words";
|
||||
|
||||
public static final List<Passage> PASSAGES = new ArrayList<>();
|
||||
public static final List<Query> QUERIES = new ArrayList<>();
|
||||
|
||||
private static void error(String code, String title, String body) {
|
||||
PASSAGES.add(new Passage(code, "error", "Error " + code + " - " + title + ". " + body));
|
||||
}
|
||||
|
||||
private static void item(String id, String kind, String text) {
|
||||
PASSAGES.add(new Passage(id, kind, text));
|
||||
}
|
||||
|
||||
private static void guide(String id, String title, String body) {
|
||||
PASSAGES.add(new Passage(id, "guide", title + ". " + body));
|
||||
}
|
||||
|
||||
private static void q(String category, String pid, String text) {
|
||||
QUERIES.add(new Query(text, pid, category));
|
||||
}
|
||||
|
||||
static {
|
||||
// ---- 40 error codes -------------------------------------------------------------------------------
|
||||
error("E-4101", "Missing idempotency key", "Requests that create a charge must send an Idempotency-Key header. A request without one is rejected before any money moves. Generate a fresh UUID for each logical payment attempt.");
|
||||
error("E-4102", "Idempotency key reused with a different body", "The same Idempotency-Key arrived with a different request body. Orbit replays the stored response only when the body is identical, and otherwise refuses the call.");
|
||||
error("E-4103", "Unsupported currency", "The currency code is not enabled for this merchant account. Enable it in the dashboard under Settings, or convert the amount before calling the API.");
|
||||
error("E-4104", "Amount below minimum", "The charge is smaller than the minimum for its currency, which is fifty cents for dollars. Orbit does not process smaller amounts.");
|
||||
error("E-4105", "Amount above limit", "The amount is higher than the per-transaction ceiling set for the merchant. Ask support to raise the ceiling or split the order into several charges.");
|
||||
error("E-4106", "Webhook signature mismatch", "The signature header on a webhook did not match the value computed with the endpoint secret. Verify against the raw request body, not the parsed JSON, because re-serialising changes the bytes.");
|
||||
error("E-4107", "Webhook timestamp too old", "The timestamp inside the webhook signature is more than five minutes in the past. This guards against replays. Check for clock drift on the receiving server.");
|
||||
error("E-4108", "Card expired", "The expiry date on the card has passed. Ask the customer for a different card or for the new expiry date before trying again.");
|
||||
error("E-4109", "Card number invalid", "The card number fails the Luhn checksum, which usually means a typing mistake. Orbit rejects it locally without contacting the issuer.");
|
||||
error("E-4110", "Strong customer authentication required", "The issuer demands a 3-D Secure challenge before it will approve the payment. Redirect the customer to the authentication page and confirm the payment afterwards.");
|
||||
error("E-4111", "Authentication challenge failed", "The customer abandoned or failed the 3-D Secure challenge. The payment cannot be completed with this card; offer another payment method.");
|
||||
error("E-4112", "Issuer declined without a reason", "The card issuer returned a generic do-not-honor response. Orbit cannot see the real cause. The customer should call their bank or use a different card.");
|
||||
error("E-4113", "Insufficient funds", "The account behind the card does not have enough available balance for the amount. A later retry may succeed once the customer tops up.");
|
||||
error("E-4114", "Refund larger than the captured amount", "The refund requested is more than what was captured on the charge, after earlier partial refunds are taken off. Lower the refund amount.");
|
||||
error("E-4115", "Refund window closed", "Refunds to cards are possible for 180 days after capture. After that the money must be returned by bank transfer outside Orbit.");
|
||||
error("E-4116", "Charge already captured", "A capture was attempted on a charge that has already been captured in full. Fetch the charge first and check its status instead of capturing blindly.");
|
||||
error("E-4117", "Authorization expired", "An authorization that is not captured within seven days lapses and the hold is released. Create a new authorization if the goods ship later.");
|
||||
error("E-4118", "Payout account not verified", "Payouts are blocked because the bank account has not passed verification. Confirm the two micro-deposits shown in the dashboard to unlock transfers.");
|
||||
error("E-4119", "Payout below threshold", "The available balance is under ten units of currency, so the scheduled payout was skipped. It will be included in the next run once the balance grows.");
|
||||
error("E-4120", "Dispute evidence deadline passed", "Evidence for a chargeback must be submitted before the deadline printed on the dispute. After it passes the dispute is lost automatically.");
|
||||
error("E-4121", "Customer not found", "No customer exists with the identifier supplied. It may belong to the other mode, since test and live objects are kept apart.");
|
||||
error("E-4122", "Payment method not attached", "The payment method is not attached to this customer. Attach it first, or pass it directly on a one-off charge.");
|
||||
error("E-4123", "Duplicate customer email", "A customer with this email address already exists and the account is configured to enforce unique emails. Reuse the existing customer.");
|
||||
error("E-4124", "Plan archived", "The subscription plan was archived and cannot be chosen by new subscribers. Existing subscribers keep it until they cancel.");
|
||||
error("E-4125", "Proration not allowed during trial", "Prorated billing cannot be applied while the subscription is still in its free trial. Change the plan at the end of the trial instead.");
|
||||
error("E-4126", "Coupon expired", "The coupon passed its redemption end date. Create a new coupon or extend the validity in the dashboard.");
|
||||
error("E-4127", "Coupon already redeemed", "This single-use coupon was already redeemed by another order. Issue a fresh code for each customer.");
|
||||
error("E-4128", "Tax identifier invalid", "The tax registration number does not match the format for its country. Check the prefix and the number of digits.");
|
||||
error("E-4129", "API key lacks permission", "The key in use does not carry the scope needed for this endpoint. Create a restricted key that includes it.");
|
||||
error("E-4130", "API key revoked", "The secret key was revoked, either by a team member or automatically after it was found in a public repository. Issue a new key and rotate it in your deployment.");
|
||||
error("E-4131", "Rate limit exceeded", "More than one hundred requests per second arrived on this key. Back off with jitter and respect the Retry-After header.");
|
||||
error("E-4132", "Request body too large", "The request payload is bigger than one megabyte. Upload large files through the file endpoint and reference them by identifier.");
|
||||
error("E-4133", "Unknown parameter", "The request contains a field that this API version does not define. Remove it or pin a newer API version.");
|
||||
error("E-4134", "Metadata key too long", "A metadata key is longer than forty characters. Shorten the key names.");
|
||||
error("E-4135", "Too many metadata keys", "An object may carry at most fifty metadata keys. Store bulky data in your own database and keep only a reference.");
|
||||
error("E-4136", "Test card used in live mode", "One of the published test card numbers was used with a live key. Switch to the test key, or use a real card.");
|
||||
error("E-4137", "Live card used in test mode", "A real card number was sent with a test key. Orbit never charges real cards in test mode, so use a test card.");
|
||||
error("E-4138", "Terminal offline", "The card terminal has not sent a heartbeat and cannot take payments. Check its network connection and power.");
|
||||
error("E-4139", "Terminal firmware outdated", "The terminal runs firmware older than the minimum supported release. Connect it to power and wifi so it can update overnight.");
|
||||
error("E-4140", "Receipt email bounced", "The receipt could not be delivered because the address rejected it. The payment itself is unaffected.");
|
||||
|
||||
// ---- 12 hardware, plan and support identifiers ------------------------------------------------------
|
||||
item("OP-TERM-5520", "sku", "OP-TERM-5520 countertop card terminal with a 4G fallback modem. Accepts chip, tap and swipe, and keeps working for up to eight hours of sales if the shop network drops.");
|
||||
item("OP-TERM-5530", "sku", "OP-TERM-5530 handheld card terminal with a built-in receipt printer and a six hour battery. Intended for table service in restaurants.");
|
||||
item("OP-TERM-6100", "sku", "OP-TERM-6100 unattended card reader for kiosks and vending machines. Vandal-resistant housing, no display, powered over the serial port.");
|
||||
item("OP-DOCK-210", "sku", "OP-DOCK-210 charging dock for the handheld terminals. Holds a single unit and charges it to full in about ninety minutes.");
|
||||
item("OP-PRNT-77", "sku", "OP-PRNT-77 thermal paper rolls, 58 millimetres wide, sold in boxes of fifty. Fits the handheld terminal.");
|
||||
item("OP-SIM-EU1", "sku", "OP-SIM-EU1 data SIM for terminals used across the European Union, with roaming included and no monthly minimum.");
|
||||
item("OP-STND-45", "sku", "OP-STND-45 adjustable stand for countertop terminals. Tilts from zero to forty-five degrees and has a lockable base.");
|
||||
item("OP-PLAN-STD", "plan", "OP-PLAN-STD standard pricing plan: 2.4 percent plus twenty cents per successful card charge, no monthly fee, payouts every two business days.");
|
||||
item("OP-PLAN-PLUS", "plan", "OP-PLAN-PLUS pricing plan for growing merchants: 1.9 percent plus twenty cents per charge, a monthly fee, and next business day payouts.");
|
||||
item("OP-PLAN-ENT", "plan", "OP-PLAN-ENT enterprise plan with negotiated rates, a named account manager and a custom payout calendar.");
|
||||
item("OP-SUP-GOLD", "support", "OP-SUP-GOLD gold support tier: a human answers within one hour at any time of day, with a phone line for outages.");
|
||||
item("OP-SUP-SILV", "support", "OP-SUP-SILV silver support tier: email answers within one business day, no phone line.");
|
||||
|
||||
// ---- 12 configuration settings ---------------------------------------------------------------------
|
||||
item("orbit.retry.max_attempts", "setting", "Setting orbit.retry.max_attempts is the number of times the client library repeats a failed request before giving up. The default is 5.");
|
||||
item("orbit.retry.backoff_ms", "setting", "Setting orbit.retry.backoff_ms is the initial wait between retries in milliseconds. It doubles after each attempt. The default is 250.");
|
||||
item("orbit.webhook.tolerance_seconds", "setting", "Setting orbit.webhook.tolerance_seconds is how old a webhook signature timestamp may be before verification fails. The default is 300.");
|
||||
item("orbit.http.timeout_ms", "setting", "Setting orbit.http.timeout_ms is the total time the client waits for a response. The default is 30000.");
|
||||
item("orbit.capture.mode", "setting", "Setting orbit.capture.mode chooses between automatic capture at authorization and manual capture later. The default is automatic.");
|
||||
item("orbit.idempotency.ttl_hours", "setting", "Setting orbit.idempotency.ttl_hours is how long Orbit remembers an idempotency key and replays its response. The default is 24.");
|
||||
item("orbit.radar.threshold", "setting", "Setting orbit.radar.threshold is the risk score from 0 to 100 above which a payment is blocked. The default is 75.");
|
||||
item("orbit.payout.schedule", "setting", "Setting orbit.payout.schedule picks daily, weekly or monthly payouts. The default is daily.");
|
||||
item("orbit.receipts.locale", "setting", "Setting orbit.receipts.locale is the language of emailed receipts. The default follows the customer's browser.");
|
||||
item("orbit.terminal.heartbeat_seconds", "setting", "Setting orbit.terminal.heartbeat_seconds is how often a terminal reports in. After three missed beats it is shown as offline. The default is 60.");
|
||||
item("orbit.log.level", "setting", "Setting orbit.log.level controls how much the client library writes to its log. The default is warn.");
|
||||
item("orbit.cache.customer_ttl_s", "setting", "Setting orbit.cache.customer_ttl_s is how long customer objects are cached inside the client. The default is 120.");
|
||||
|
||||
// ---- 30 guides in plain language --------------------------------------------------------------------
|
||||
guide("G-01", "Recovering failed subscription renewals", "When a renewal charge fails, Orbit runs a dunning schedule: it retries after one, three and five days, then marks the subscription past due and emails the customer a link to update their card.");
|
||||
guide("G-02", "Paying several sellers from one order", "A marketplace can divide the proceeds of a single charge between sellers using connected accounts and transfers. Each seller is paid out on their own schedule and the platform keeps its fee.");
|
||||
guide("G-03", "Authorise now, capture on dispatch", "Place a hold on the card when the order is made and capture it when the goods ship. The hold guarantees the funds for seven days and can be released without any fee.");
|
||||
guide("G-04", "Responding to a chargeback", "When a cardholder disputes a payment their bank takes the money back at once. Gather proof of delivery and customer communication and submit it through the dispute page before the deadline.");
|
||||
guide("G-05", "Trying the integration without real money", "Test mode mirrors live mode with separate keys and published card numbers that simulate approvals, declines and challenges. Nothing in test mode moves real funds.");
|
||||
guide("G-06", "Remembering a shopper's card", "Attach a payment method to a customer object after the first purchase. Later orders can then reuse it, and the shopper never types the number again.");
|
||||
guide("G-07", "Taking payment from people without cards", "Orbit supports bank redirects, wallets and instant bank transfers. Enable the methods in the dashboard and they appear in the checkout automatically for eligible countries.");
|
||||
guide("G-08", "When refunded money arrives", "A refund is sent to the bank immediately but most issuers need five to ten business days to show it on the customer's statement.");
|
||||
guide("G-09", "What shows on the cardholder's statement", "The statement descriptor is the short text a cardholder sees next to the charge. Choose one customers will recognise to avoid unnecessary disputes.");
|
||||
guide("G-10", "Blocking suspicious orders", "Radar scores every payment for fraud risk using signals across the network. Rules let you block, review or allow payments based on the score, the country or the card type.");
|
||||
guide("G-11", "When funds reach your bank account", "Settled funds are paid to your bank on a rolling schedule. New accounts wait longer for the first payout while Orbit builds a history.");
|
||||
guide("G-12", "If your server is unreachable", "When your endpoint does not answer, Orbit redelivers the event with increasing gaps for up to three days. You can also list missed events through the events endpoint and replay them.");
|
||||
guide("G-13", "Charging in one currency and settling in another", "Present prices in the shopper's currency and Orbit converts to your settlement currency at the daily rate, adding a small conversion fee to the payout.");
|
||||
guide("G-14", "Adding sales tax to bills", "Orbit Tax calculates the rate from the customer's address and the product category, adds it to the invoice and keeps the reports you need for filing.");
|
||||
guide("G-15", "Giving new customers a free period", "Set a trial length on the plan. The customer is not charged until the trial ends, and a reminder is emailed three days before the first payment.");
|
||||
guide("G-16", "Putting a subscription on hold", "Pausing stops invoices while keeping the customer and their payment method. Billing resumes on the date you choose and the cycle is shifted accordingly.");
|
||||
guide("G-17", "Handing transactions to the accountants", "Export every transaction, fee and payout for a date range as a spreadsheet. The report reconciles to the bank deposits line by line.");
|
||||
guide("G-18", "Restricting what a team member can see", "Dashboard roles control who can view balances, issue refunds or read customer details. Give support staff the limited role so they cannot see payout accounts.");
|
||||
guide("G-19", "Keeping European data in Europe", "Choose the EU region when you create the account and customer and card data are stored and processed only in European data centres.");
|
||||
guide("G-20", "Staying out of card data rules", "Hosted fields and the checkout page collect card numbers on Orbit's servers, so they never touch yours and your compliance questionnaire stays short.");
|
||||
guide("G-21", "Accepting taps on a phone", "The mobile SDK turns a supported phone into a contactless reader, with no extra hardware, for cards and wallets.");
|
||||
guide("G-22", "Customising the emailed receipt", "Add your logo, brand colour and a footer message to receipts under Settings. Receipts are sent automatically after each successful payment.");
|
||||
guide("G-23", "Resending a request that timed out", "If a call times out you cannot know whether it succeeded. Send it again with the same idempotency key and Orbit returns the original result instead of creating a second charge.");
|
||||
guide("G-24", "Taking part of an authorised amount", "You may capture less than was authorised, for example when one item of an order is out of stock. The rest of the hold is released straight away.");
|
||||
guide("G-25", "Letting customers add a tip", "Collect a gratuity on the terminal after the amount is shown. The tip is added to the same charge and reported separately.");
|
||||
guide("G-26", "Showing local prices", "Define a price per currency on each product so shoppers see round local amounts instead of converted ones.");
|
||||
guide("G-27", "Clearing out sandbox data", "Reset the test environment from the developer settings page to delete test customers, charges and subscriptions and start clean.");
|
||||
guide("G-28", "Seeing who changed what", "The audit log records every dashboard action with the person, time and address, and can be exported for your security team.");
|
||||
guide("G-29", "Signing in with your company directory", "Single sign-on lets the team log in to the dashboard through your identity provider, and access is removed automatically when someone leaves.");
|
||||
guide("G-30", "Checking whether Orbit is having problems", "The status page lists current incidents and past outages per region, and you can subscribe to updates by email or webhook.");
|
||||
|
||||
// ---- 60 queries --------------------------------------------------------------------------------------
|
||||
// 12 where the only useful words are an error code
|
||||
q(EXACT_CODE, "E-4106", "What does error E-4106 mean?");
|
||||
q(EXACT_CODE, "E-4116", "Why am I getting E-4116?");
|
||||
q(EXACT_CODE, "E-4118", "We see E-4118 in the logs, what should we do?");
|
||||
q(EXACT_CODE, "E-4120", "Explain E-4120");
|
||||
q(EXACT_CODE, "E-4122", "Support told me to look up E-4122");
|
||||
q(EXACT_CODE, "E-4126", "What is E-4126?");
|
||||
q(EXACT_CODE, "E-4129", "How do I fix E-4129?");
|
||||
q(EXACT_CODE, "E-4131", "Our service started returning E-4131 this morning");
|
||||
q(EXACT_CODE, "E-4133", "Meaning of E-4133");
|
||||
q(EXACT_CODE, "E-4135", "Is E-4135 our fault?");
|
||||
q(EXACT_CODE, "E-4138", "What causes E-4138 on the shop floor?");
|
||||
q(EXACT_CODE, "E-4140", "Customer reports E-4140, is the money safe?");
|
||||
|
||||
// 12 where the identifier is a SKU, plan or setting name
|
||||
q(EXACT_ID, "OP-TERM-5520", "Does OP-TERM-5520 work when the internet is down?");
|
||||
q(EXACT_ID, "OP-TERM-5530", "Tell me about OP-TERM-5530");
|
||||
q(EXACT_ID, "OP-TERM-6100", "Can OP-TERM-6100 be used in a vending machine?");
|
||||
q(EXACT_ID, "OP-DOCK-210", "How long does OP-DOCK-210 take to charge a unit?");
|
||||
q(EXACT_ID, "OP-PRNT-77", "What size is OP-PRNT-77?");
|
||||
q(EXACT_ID, "OP-PLAN-PLUS", "What are the rates on OP-PLAN-PLUS?");
|
||||
q(EXACT_ID, "OP-SUP-GOLD", "Response time for OP-SUP-GOLD");
|
||||
q(EXACT_ID, "orbit.retry.backoff_ms", "What is the default for orbit.retry.backoff_ms?");
|
||||
q(EXACT_ID, "orbit.radar.threshold", "Value of orbit.radar.threshold");
|
||||
q(EXACT_ID, "orbit.idempotency.ttl_hours", "What does orbit.idempotency.ttl_hours do?");
|
||||
q(EXACT_ID, "orbit.terminal.heartbeat_seconds", "Default of orbit.terminal.heartbeat_seconds");
|
||||
q(EXACT_ID, "orbit.cache.customer_ttl_s", "Explain orbit.cache.customer_ttl_s");
|
||||
|
||||
// 24 plain-language questions written without the passage's own distinctive words
|
||||
q(PARAPHRASE, "G-01", "What happens when a customer's monthly payment does not go through?");
|
||||
q(PARAPHRASE, "G-02", "How can an online mall split one shopper's purchase among several vendors?");
|
||||
q(PARAPHRASE, "G-03", "Can I reserve money on a card and only take it once the parcel leaves?");
|
||||
q(PARAPHRASE, "G-04", "A buyer told their bank they never got the item. What now?");
|
||||
q(PARAPHRASE, "G-05", "How do I experiment safely before launching?");
|
||||
q(PARAPHRASE, "G-06", "Store a shopper's card so they do not re-enter it next time");
|
||||
q(PARAPHRASE, "G-07", "Accept money from people who do not own a credit card");
|
||||
q(PARAPHRASE, "G-08", "How long until the buyer sees their money back?");
|
||||
q(PARAPHRASE, "G-09", "Which words appear on the customer's bank statement?");
|
||||
q(PARAPHRASE, "G-10", "Stop dodgy purchases automatically");
|
||||
q(PARAPHRASE, "G-11", "When will the money land in our account?");
|
||||
q(PARAPHRASE, "G-12", "What if our website is down when you send us news?");
|
||||
q(PARAPHRASE, "G-13", "Show prices in euros but receive dollars");
|
||||
q(PARAPHRASE, "G-14", "Include VAT on our bills without doing sums");
|
||||
q(PARAPHRASE, "G-15", "Let new users try the product for two weeks at no cost");
|
||||
q(PARAPHRASE, "G-16", "A client wants to stop being billed for the summer");
|
||||
q(PARAPHRASE, "G-17", "Give our bookkeeper everything that came in last quarter");
|
||||
q(PARAPHRASE, "G-18", "Limit what a helpdesk employee can look at");
|
||||
q(PARAPHRASE, "G-19", "Our lawyers want shopper information to stay on this continent");
|
||||
q(PARAPHRASE, "G-20", "How do we avoid ever handling raw card numbers ourselves?");
|
||||
q(PARAPHRASE, "G-21", "Take contactless payments with just a smartphone");
|
||||
q(PARAPHRASE, "G-23", "Our call hung and we do not know if it worked. Is it safe to send it again?");
|
||||
q(PARAPHRASE, "G-24", "One item in the basket is sold out, can we bill only for the rest?");
|
||||
q(PARAPHRASE, "G-28", "Which colleague altered the settings last Tuesday?");
|
||||
|
||||
// 12 that carry a code and also say what it is about
|
||||
q(MIXED, "E-4106", "E-4106 keeps appearing in our Express webhook handler, is the body being parsed?");
|
||||
q(MIXED, "E-4113", "Customer card gets E-4113 although they swear there is money on it");
|
||||
q(MIXED, "E-4117", "E-4117 on capture, the goods shipped ten days after the order");
|
||||
q(MIXED, "E-4125", "Plan change during trial returns E-4125");
|
||||
q(MIXED, "E-4130", "E-4130 on every call since the key leaked into our repository");
|
||||
q(MIXED, "E-4115", "Cannot refund a purchase from last year, E-4115");
|
||||
q(MIXED, "E-4110", "E-4110 and the customer is asked for a bank challenge");
|
||||
q(MIXED, "E-4127", "Promo code gives E-4127 for the second shopper");
|
||||
q(MIXED, "E-4139", "Terminal says E-4139 and needs an update");
|
||||
q(MIXED, "E-4101", "Getting E-4101 on the charge endpoint, our client does not set the idempotency header");
|
||||
q(MIXED, "E-4131", "E-4131 during the flash sale, too many requests per second");
|
||||
q(MIXED, "E-4119", "Why was our payout skipped, E-4119, the balance was small");
|
||||
}
|
||||
|
||||
private KnowledgeBase() {
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,64 @@
|
||||
package com.ankurm.rerank;
|
||||
|
||||
import java.io.IOException;
|
||||
import java.io.InputStream;
|
||||
import java.net.URI;
|
||||
import java.nio.file.Files;
|
||||
import java.nio.file.Path;
|
||||
import java.nio.file.StandardCopyOption;
|
||||
|
||||
/**
|
||||
* Where the two models come from. Both are public Hugging Face repositories; the first run downloads about
|
||||
* 200 MB into {@code ~/.cache/rerank} and every later run is offline. No API key is involved.
|
||||
* <ul>
|
||||
* <li>bi-encoder: sentence-transformers/all-MiniLM-L6-v2 (384-dim embeddings, retrieval stage), run by {@link BiEncoder}</li>
|
||||
* <li>cross-encoder: cross-encoder/ms-marco-MiniLM-L-6-v2 (reranking stage), fp32 and an int8 variant</li>
|
||||
* </ul>
|
||||
*/
|
||||
public final class Models {
|
||||
|
||||
private static final String BI = "https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2/resolve/main/";
|
||||
private static final String CROSS = "https://huggingface.co/cross-encoder/ms-marco-MiniLM-L-6-v2/resolve/main/";
|
||||
private static final Path CACHE = Path.of(System.getProperty("user.home"), ".cache", "rerank");
|
||||
|
||||
public static final String CROSS_FP32 = "onnx/model.onnx";
|
||||
public static final String CROSS_INT8 = "onnx/model_quint8_avx2.onnx";
|
||||
|
||||
private static BiEncoder bi;
|
||||
|
||||
private Models() {
|
||||
}
|
||||
|
||||
public static synchronized BiEncoder biEncoder() {
|
||||
if (bi == null) {
|
||||
bi = new BiEncoder(fetch(BI + "onnx/model.onnx", CACHE.resolve("bi-encoder/model.onnx")),
|
||||
fetch(BI + "tokenizer.json", CACHE.resolve("bi-encoder/tokenizer.json")), 256);
|
||||
}
|
||||
return bi;
|
||||
}
|
||||
|
||||
public static Path crossTokenizer() {
|
||||
return fetch(CROSS + "tokenizer.json", CACHE.resolve("cross-encoder/tokenizer.json"));
|
||||
}
|
||||
|
||||
public static Path crossModel(String file) {
|
||||
return fetch(CROSS + file, CACHE.resolve("cross-encoder").resolve(file));
|
||||
}
|
||||
|
||||
private static Path fetch(String url, Path target) {
|
||||
if (Files.exists(target)) {
|
||||
return target;
|
||||
}
|
||||
try {
|
||||
Files.createDirectories(target.getParent());
|
||||
Path tmp = target.resolveSibling(target.getFileName() + ".part");
|
||||
try (InputStream in = URI.create(url).toURL().openStream()) {
|
||||
Files.copy(in, tmp, StandardCopyOption.REPLACE_EXISTING);
|
||||
}
|
||||
Files.move(tmp, target, StandardCopyOption.REPLACE_EXISTING);
|
||||
return target;
|
||||
} catch (IOException e) {
|
||||
throw new IllegalStateException("could not download " + url, e);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,49 @@
|
||||
package com.ankurm.rerank;
|
||||
|
||||
import java.util.ArrayList;
|
||||
import java.util.Comparator;
|
||||
import java.util.HashMap;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
|
||||
import org.springframework.ai.document.Document;
|
||||
import org.springframework.ai.rag.Query;
|
||||
import org.springframework.ai.rag.postretrieval.document.DocumentPostProcessor;
|
||||
|
||||
/**
|
||||
* Spring AI's hook between "documents retrieved" and "documents put in the prompt". Hand it to
|
||||
* {@code RetrievalAugmentationAdvisor.builder().documentPostProcessors(...)} and the model sees the
|
||||
* cross-encoder's top N instead of the retriever's.
|
||||
*/
|
||||
public class RerankingPostProcessor implements DocumentPostProcessor {
|
||||
|
||||
private final CrossEncoder crossEncoder;
|
||||
private final int topN;
|
||||
|
||||
public RerankingPostProcessor(CrossEncoder crossEncoder, int topN) {
|
||||
this.crossEncoder = crossEncoder;
|
||||
this.topN = topN;
|
||||
}
|
||||
|
||||
@Override
|
||||
public List<Document> process(Query query, List<Document> documents) {
|
||||
if (documents.isEmpty()) {
|
||||
return documents;
|
||||
}
|
||||
float[] scores = crossEncoder.score(query.text(), documents.stream().map(Document::getText).toList());
|
||||
List<Integer> order = new ArrayList<>();
|
||||
for (int i = 0; i < scores.length; i++) {
|
||||
order.add(i);
|
||||
}
|
||||
order.sort(Comparator.comparingDouble((Integer i) -> -scores[i]));
|
||||
List<Document> out = new ArrayList<>();
|
||||
for (int i : order.subList(0, Math.min(topN, order.size()))) {
|
||||
Document d = documents.get(i);
|
||||
Map<String, Object> meta = new HashMap<>(d.getMetadata());
|
||||
meta.put("retrieval_score", d.getScore());
|
||||
meta.put("rerank_score", scores[i]);
|
||||
out.add(d.mutate().metadata(meta).score((double) scores[i]).build());
|
||||
}
|
||||
return out;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,110 @@
|
||||
package com.ankurm.rerank;
|
||||
|
||||
import static org.junit.jupiter.api.Assertions.assertTrue;
|
||||
|
||||
import java.util.ArrayList;
|
||||
import java.util.Collections;
|
||||
import java.util.List;
|
||||
|
||||
import com.ankurm.rerank.support.Fixture;
|
||||
import com.ankurm.rerank.support.Metrics;
|
||||
import com.ankurm.rerank.support.Transcript;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.springframework.ai.document.Document;
|
||||
import org.springframework.ai.rag.Query;
|
||||
|
||||
/**
|
||||
* Wall-clock cost of the second stage. These numbers come from the machine that ran the test (see the
|
||||
* header of each file) and will differ on yours; the shape of the curves is what carries over.
|
||||
*/
|
||||
class LatencyTest {
|
||||
|
||||
private static final int WARMUP = 3;
|
||||
private static final int RUNS = 12;
|
||||
|
||||
private static double[] timeRerank(CrossEncoder ce, int pool) {
|
||||
List<Double> ms = new ArrayList<>();
|
||||
for (int i = 0; i < WARMUP + RUNS; i++) {
|
||||
KnowledgeBase.Query q = KnowledgeBase.QUERIES.get((i * 7) % KnowledgeBase.QUERIES.size());
|
||||
List<String> texts = Fixture.retriever().withTopK(pool).retrieve(new Query(q.text())).stream().map(Document::getText).toList();
|
||||
long t0 = System.nanoTime();
|
||||
ce.score(q.text(), texts);
|
||||
double elapsed = (System.nanoTime() - t0) / 1e6;
|
||||
if (i >= WARMUP) {
|
||||
ms.add(elapsed);
|
||||
}
|
||||
}
|
||||
Collections.sort(ms);
|
||||
return new double[] {ms.get(ms.size() / 2), ms.get(ms.size() - 1)};
|
||||
}
|
||||
|
||||
private static String machine() {
|
||||
return Runtime.getRuntime().availableProcessors() + " vCPU, JDK " + System.getProperty("java.version");
|
||||
}
|
||||
|
||||
@Test
|
||||
void candidatesAndBatchSize() {
|
||||
Models.crossModel(Models.CROSS_FP32);
|
||||
try (Transcript t = new Transcript("06-latency-by-candidates-and-batch.txt", "Time to rerank N candidates (median and slowest of " + RUNS + " calls, after " + WARMUP + " warm-up calls)");
|
||||
CrossEncoder b1 = new CrossEncoder(Models.crossModel(Models.CROSS_FP32), Models.crossTokenizer(), Fixture.MAX_LEN, 1, 0);
|
||||
CrossEncoder b8 = new CrossEncoder(Models.crossModel(Models.CROSS_FP32), Models.crossTokenizer(), Fixture.MAX_LEN, 8, 0);
|
||||
CrossEncoder b32 = new CrossEncoder(Models.crossModel(Models.CROSS_FP32), Models.crossTokenizer(), Fixture.MAX_LEN, 32, 0);
|
||||
CrossEncoder one = new CrossEncoder(Models.crossModel(Models.CROSS_FP32), Models.crossTokenizer(), Fixture.MAX_LEN, 32, 1)) {
|
||||
t.line("machine: %s", machine());
|
||||
t.line("fp32 cross-encoder, max length %d tokens; the first column of each cell is the median in ms, the second the slowest", Fixture.MAX_LEN);
|
||||
t.blank();
|
||||
t.line("%-12s %-14s %-14s %-14s %-14s", "candidates", "batch 1", "batch 8", "batch 32", "batch 32, 1 thread");
|
||||
for (int pool : new int[] {10, 20, 50, 94}) {
|
||||
double[] a = timeRerank(b1, pool);
|
||||
double[] b = timeRerank(b8, pool);
|
||||
double[] c = timeRerank(b32, pool);
|
||||
double[] d = timeRerank(one, pool);
|
||||
t.line("%-12d %-14s %-14s %-14s %-14s", pool, fmt(a), fmt(b), fmt(c), fmt(d));
|
||||
assertTrue(a[0] > 0 && d[0] > 0);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@Test
|
||||
void int8AndTheWholePipeline() {
|
||||
CrossEncoder fp32 = Fixture.fp32();
|
||||
CrossEncoder int8 = Fixture.int8();
|
||||
try (Transcript t = new Transcript("07-latency-against-quality.txt", "Quality and added latency for each setting (60 questions; latency is the median of " + RUNS + " calls)")) {
|
||||
t.line("machine: %s", machine());
|
||||
t.line("first stage (embed the question, compare with 94 vectors): median %.1f ms", timeFirstStage());
|
||||
t.blank();
|
||||
t.line("%-28s %-8s %-8s %-8s %-8s %s", "setting", "hit@1", "hit@3", "hit@5", "MRR@10", "rerank ms (median)");
|
||||
Metrics.Score bi = Metrics.score(KnowledgeBase.QUERIES, Fixture::biRanking);
|
||||
t.line("%-28s %-8s %-8s %-8s %-8.3f %s", "no reranker", bi.hit1() + "/60", bi.hit3() + "/60", bi.hit5() + "/60", bi.mrr(), "0");
|
||||
for (int pool : new int[] {5, 10, 20, 50}) {
|
||||
for (boolean q8 : new boolean[] {false, true}) {
|
||||
CrossEncoder ce = q8 ? int8 : fp32;
|
||||
Metrics.Score s = Metrics.score(KnowledgeBase.QUERIES, q -> Fixture.reranked(ce, q, pool));
|
||||
double[] lat = timeRerank(ce, pool);
|
||||
t.line("%-28s %-8s %-8s %-8s %-8.3f %.0f", (q8 ? "int8" : "fp32") + ", top " + pool, s.hit1() + "/60", s.hit3() + "/60", s.hit5() + "/60", s.mrr(), lat[0]);
|
||||
}
|
||||
}
|
||||
long agree = KnowledgeBase.QUERIES.stream().filter(q -> Fixture.reranked(fp32, q, 50).get(0).equals(Fixture.reranked(int8, q, 50).get(0))).count();
|
||||
t.blank();
|
||||
t.line("int8 and fp32 put the same passage first for %d of %d questions (top 50)", agree, KnowledgeBase.QUERIES.size());
|
||||
}
|
||||
}
|
||||
|
||||
private static double timeFirstStage() {
|
||||
List<Double> ms = new ArrayList<>();
|
||||
for (int i = 0; i < WARMUP + RUNS; i++) {
|
||||
String q = KnowledgeBase.QUERIES.get((i * 7) % KnowledgeBase.QUERIES.size()).text();
|
||||
long t0 = System.nanoTime();
|
||||
Fixture.retriever().withTopK(50).retrieve(new Query(q));
|
||||
if (i >= WARMUP) {
|
||||
ms.add((System.nanoTime() - t0) / 1e6);
|
||||
}
|
||||
}
|
||||
Collections.sort(ms);
|
||||
return ms.get(ms.size() / 2);
|
||||
}
|
||||
|
||||
private static String fmt(double[] v) {
|
||||
return String.format("%.0f / %.0f", v[0], v[1]);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,59 @@
|
||||
package com.ankurm.rerank;
|
||||
|
||||
import static org.junit.jupiter.api.Assertions.assertEquals;
|
||||
import static org.junit.jupiter.api.Assertions.assertTrue;
|
||||
|
||||
import java.util.Arrays;
|
||||
import java.util.List;
|
||||
|
||||
import ai.djl.huggingface.tokenizers.Encoding;
|
||||
import com.ankurm.rerank.support.Fixture;
|
||||
import com.ankurm.rerank.support.Transcript;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.springframework.ai.document.Document;
|
||||
import org.springframework.ai.rag.Query;
|
||||
|
||||
/** What goes into a cross-encoder, and one question reranked by hand. */
|
||||
class PairEncodingTest {
|
||||
|
||||
@Test
|
||||
void whatTheModelSees() {
|
||||
CrossEncoder ce = Fixture.fp32();
|
||||
try (Transcript t = new Transcript("01-pair-encoding.txt", "A (query, passage) pair as the cross-encoder receives it")) {
|
||||
t.line("model inputs: %s", ce.inputNames());
|
||||
Encoding e = ce.encode("What is E-4126?", "Error E-4126 - Refund window closed. Refunds can be issued for 120 days.");
|
||||
t.line("tokens : %s", Arrays.toString(e.getTokens()));
|
||||
t.line("token_type_ids: %s", Arrays.toString(e.getTypeIds()));
|
||||
t.line("attention_mask: %s", Arrays.toString(e.getAttentionMask()));
|
||||
assertEquals("[CLS]", e.getTokens()[0]);
|
||||
assertTrue(Arrays.stream(e.getTypeIds()).anyMatch(x -> x == 1));
|
||||
}
|
||||
}
|
||||
|
||||
@Test
|
||||
void oneQuestionByHand() {
|
||||
String question = "A buyer told their bank they never got the item. What now?";
|
||||
List<Document> candidates = Fixture.retriever().withTopK(8).retrieve(new Query(question));
|
||||
float[] scores = Fixture.fp32().score(question, candidates.stream().map(Document::getText).toList());
|
||||
try (Transcript t = new Transcript("02-one-question-by-hand.txt", "One question, the bi-encoder's top 8, rescored")) {
|
||||
t.line("question: %s", question);
|
||||
t.line("gold passage: G-04");
|
||||
t.blank();
|
||||
t.line("%-9s %-7s %-12s %s", "retrieved", "id", "cosine", "cross-encoder logit");
|
||||
int goldRetrieved = -1;
|
||||
int best = 0;
|
||||
for (int i = 0; i < candidates.size(); i++) {
|
||||
Document d = candidates.get(i);
|
||||
t.line("%-9d %-7s %-12.3f %.2f", i + 1, d.getMetadata().get("pid"), d.getScore(), scores[i]);
|
||||
if ("G-04".equals(d.getMetadata().get("pid"))) {
|
||||
goldRetrieved = i + 1;
|
||||
}
|
||||
if (scores[i] > scores[best]) {
|
||||
best = i;
|
||||
}
|
||||
}
|
||||
t.blank();
|
||||
t.line("bi-encoder rank of G-04: %d; cross-encoder winner: %s", goldRetrieved, candidates.get(best).getMetadata().get("pid"));
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,63 @@
|
||||
package com.ankurm.rerank;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
|
||||
import java.util.List;
|
||||
|
||||
import com.ankurm.rerank.support.Fixture;
|
||||
import com.ankurm.rerank.support.RecordingChatModel;
|
||||
import com.ankurm.rerank.support.Transcript;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.springframework.ai.chat.client.ChatClient;
|
||||
import org.springframework.ai.document.Document;
|
||||
import org.springframework.ai.rag.Query;
|
||||
import org.springframework.ai.rag.advisor.RetrievalAugmentationAdvisor;
|
||||
|
||||
/** The post-processor inside Spring AI's RetrievalAugmentationAdvisor, and what reaches the prompt. */
|
||||
class PostProcessorTest {
|
||||
|
||||
@Test
|
||||
void fiftyInFiveOut() {
|
||||
BiEncoderRetriever retriever = Fixture.retriever().withTopK(50);
|
||||
RerankingPostProcessor reranker = new RerankingPostProcessor(Fixture.fp32(), 5);
|
||||
String question = "How do I fix E-4129?";
|
||||
|
||||
try (Transcript t = new Transcript("08-post-processor-in-the-advisor.txt", "RerankingPostProcessor inside RetrievalAugmentationAdvisor (retrieve 50, keep 5)")) {
|
||||
List<Document> retrieved = retriever.retrieve(new Query(question));
|
||||
List<Document> kept = reranker.process(new Query(question), retrieved);
|
||||
t.line("question: %s", question);
|
||||
t.line("retriever returned %d documents; the post-processor kept %d", retrieved.size(), kept.size());
|
||||
t.blank();
|
||||
t.line("%-5s %-8s %-15s %-16s %s", "kept", "id", "retrieval rank", "rerank score", "text");
|
||||
for (int i = 0; i < kept.size(); i++) {
|
||||
Document d = kept.get(i);
|
||||
t.line("%-5d %-8s %-15s %-16.2f %s", i + 1, d.getMetadata().get("pid"), d.getMetadata().get("retrieval_rank"),
|
||||
(Float) d.getMetadata().get("rerank_score"), d.getText().substring(0, 48));
|
||||
}
|
||||
assertThat(retrieved).hasSize(50);
|
||||
assertThat(kept).hasSize(5);
|
||||
assertThat(kept.getFirst().getMetadata().get("pid")).isEqualTo("E-4129");
|
||||
|
||||
RecordingChatModel chat = new RecordingChatModel();
|
||||
ChatClient client = ChatClient.builder(chat).defaultAdvisors(RetrievalAugmentationAdvisor.builder()
|
||||
.documentRetriever(retriever).documentPostProcessors(reranker).build()).build();
|
||||
client.prompt().user(question).call().content();
|
||||
String prompt = chat.prompts().getFirst().getContents();
|
||||
t.section("the prompt the model received");
|
||||
t.line(prompt);
|
||||
assertThat(prompt).contains("Error E-4129");
|
||||
}
|
||||
}
|
||||
|
||||
@Test
|
||||
void edgeCases() {
|
||||
RerankingPostProcessor reranker = new RerankingPostProcessor(Fixture.fp32(), 5);
|
||||
try (Transcript t = new Transcript("09-post-processor-edge-cases.txt", "RerankingPostProcessor edge cases")) {
|
||||
List<Document> none = reranker.process(new Query("anything"), List.of());
|
||||
t.line("no documents in -> %d out", none.size());
|
||||
List<Document> three = Fixture.retriever().withTopK(3).retrieve(new Query("Why am I getting E-4116?"));
|
||||
t.line("3 documents in, topN = 5 -> %d out (nothing invented, nothing dropped)", reranker.process(new Query("Why am I getting E-4116?"), three).size());
|
||||
assertThat(none).isEmpty();
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,78 @@
|
||||
package com.ankurm.rerank;
|
||||
|
||||
import static org.junit.jupiter.api.Assertions.assertTrue;
|
||||
|
||||
import java.util.List;
|
||||
|
||||
import com.ankurm.rerank.support.Fixture;
|
||||
import com.ankurm.rerank.support.Metrics;
|
||||
import com.ankurm.rerank.support.Transcript;
|
||||
import org.junit.jupiter.api.Test;
|
||||
|
||||
/** Does reranking put the right passage nearer the top? 60 labelled questions, one gold passage each. */
|
||||
class QualityTest {
|
||||
|
||||
private static final int[] POOLS = {5, 10, 20, 50, 94};
|
||||
|
||||
@Test
|
||||
void poolSizeAgainstQuality() {
|
||||
CrossEncoder ce = Fixture.fp32();
|
||||
try (Transcript t = new Transcript("03-quality-by-pool-size.txt", "Rerank the bi-encoder's top N (60 questions, 94 passages, fp32 cross-encoder)")) {
|
||||
t.line("%-34s %-8s %-8s %-8s %s", "ranking", "hit@1", "hit@3", "hit@5", "MRR@10");
|
||||
Metrics.Score bi = Metrics.score(KnowledgeBase.QUERIES, Fixture::biRanking);
|
||||
t.line("%-34s %s", "bi-encoder alone", bi.cells());
|
||||
Metrics.Score last = null;
|
||||
for (int pool : POOLS) {
|
||||
last = Metrics.score(KnowledgeBase.QUERIES, q -> Fixture.reranked(ce, q, pool));
|
||||
t.line("%-34s %s", "bi-encoder top " + pool + " + cross-encoder", last.cells());
|
||||
}
|
||||
t.blank();
|
||||
t.line("recall of the first stage (is the gold passage inside the pool at all?):");
|
||||
for (int pool : POOLS) {
|
||||
long in = KnowledgeBase.QUERIES.stream().filter(q -> Metrics.rankOf(Fixture.biRanking(q), q.pid()) <= pool && Metrics.rankOf(Fixture.biRanking(q), q.pid()) > 0).count();
|
||||
t.line(" top %-3d: %2d/%d", pool, in, KnowledgeBase.QUERIES.size());
|
||||
}
|
||||
assertTrue(last.hit1() >= 0);
|
||||
}
|
||||
}
|
||||
|
||||
@Test
|
||||
void byQuestionKind() {
|
||||
CrossEncoder ce = Fixture.fp32();
|
||||
try (Transcript t = new Transcript("04-quality-by-question-kind.txt", "Bi-encoder alone against bi-encoder top 50 + cross-encoder, by kind of question")) {
|
||||
t.line("%-22s %-10s %-8s %-8s %-8s %s", "kind", "ranking", "hit@1", "hit@3", "hit@5", "MRR@10");
|
||||
for (String kind : List.of(KnowledgeBase.EXACT_CODE, KnowledgeBase.EXACT_ID, KnowledgeBase.PARAPHRASE, KnowledgeBase.MIXED)) {
|
||||
Metrics.Score bi = Metrics.score(Metrics.category(kind), Fixture::biRanking);
|
||||
Metrics.Score rr = Metrics.score(Metrics.category(kind), q -> Fixture.reranked(ce, q, 50));
|
||||
t.line("%-22s %-10s %s", kind, "bi-encoder", bi.cells());
|
||||
t.line("%-22s %-10s %s", "", "reranked", rr.cells());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@Test
|
||||
void whoWinsAndWhoLoses() {
|
||||
CrossEncoder ce = Fixture.fp32();
|
||||
try (Transcript t = new Transcript("05-helped-and-hurt.txt", "Gold passage rank before and after reranking the top 50 (only questions whose rank changed)")) {
|
||||
int up = 0, down = 0, same = 0;
|
||||
t.line("%-6s %-34s %-11s %s", "gold", "question", "bi-encoder", "reranked");
|
||||
for (KnowledgeBase.Query q : KnowledgeBase.QUERIES) {
|
||||
int before = Metrics.rankOf(Fixture.biRanking(q), q.pid());
|
||||
int after = Metrics.rankOf(Fixture.reranked(ce, q, 50), q.pid());
|
||||
if (after < before) {
|
||||
up++;
|
||||
} else if (after > before) {
|
||||
down++;
|
||||
} else {
|
||||
same++;
|
||||
}
|
||||
if (after != before) {
|
||||
String text = q.text().length() > 32 ? q.text().substring(0, 31) + "~" : q.text();
|
||||
t.line("%-6s %-34s %-11d %d %s", q.pid().length() > 6 ? q.pid().substring(0, 6) : q.pid(), text, before, after, after < before ? "better" : "worse");
|
||||
}
|
||||
}
|
||||
t.blank();
|
||||
t.line("better: %d, worse: %d, unchanged: %d", up, down, same);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,153 @@
|
||||
package com.ankurm.rerank;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
import static org.junit.jupiter.api.Assertions.assertThrows;
|
||||
|
||||
import java.util.ArrayList;
|
||||
import java.util.Collections;
|
||||
import java.util.List;
|
||||
|
||||
import com.ankurm.rerank.support.Fixture;
|
||||
import com.ankurm.rerank.support.Metrics;
|
||||
import com.ankurm.rerank.support.Transcript;
|
||||
import org.junit.jupiter.api.Test;
|
||||
|
||||
/** The ways a reranker goes quietly wrong. Each test reproduces one. */
|
||||
class TrapsTest {
|
||||
|
||||
@Test
|
||||
void missingTokenTypeIds() {
|
||||
try (Transcript t = new Transcript("10-trap-missing-input.txt", "Leaving out token_type_ids");
|
||||
CrossEncoder broken = new CrossEncoder(Models.crossModel(Models.CROSS_FP32), Models.crossTokenizer(), Fixture.MAX_LEN, 32, 0, false)) {
|
||||
IllegalStateException e = assertThrows(IllegalStateException.class, () -> broken.score("q", List.of("p")));
|
||||
t.line("model inputs : %s", broken.inputNames());
|
||||
t.line("we sent : [input_ids, attention_mask]");
|
||||
t.line("result : %s", e.getMessage());
|
||||
assertThat(e.getMessage()).contains("token_type_ids");
|
||||
}
|
||||
}
|
||||
|
||||
@Test
|
||||
void wrongOrderOfThePair() {
|
||||
CrossEncoder ce = Fixture.fp32();
|
||||
try (Transcript t = new Transcript("11-trap-pair-order.txt", "(passage, query) instead of (query, passage); top 20 reranked")) {
|
||||
int pool = 20;
|
||||
int good = 0, swapped = 0;
|
||||
double goodMrr = 0, swappedMrr = 0;
|
||||
for (KnowledgeBase.Query q : KnowledgeBase.QUERIES) {
|
||||
List<String> bi = Fixture.biRanking(q).subList(0, pool);
|
||||
List<String> texts = bi.stream().map(pid -> KnowledgeBase.PASSAGES.get(Fixture.pidIndex(pid)).text()).toList();
|
||||
float[] right = ce.score(q.text(), texts);
|
||||
float[] wrong = new float[texts.size()];
|
||||
for (int i = 0; i < texts.size(); i++) {
|
||||
wrong[i] = ce.score(texts.get(i), List.of(q.text()))[0];
|
||||
}
|
||||
int rr = rankAfter(bi, right, q.pid());
|
||||
int rw = rankAfter(bi, wrong, q.pid());
|
||||
good += rr == 1 ? 1 : 0;
|
||||
swapped += rw == 1 ? 1 : 0;
|
||||
goodMrr += rr > 0 ? 1.0 / rr : 0;
|
||||
swappedMrr += rw > 0 ? 1.0 / rw : 0;
|
||||
}
|
||||
t.line("%-24s hit@1 %d/60 MRR %.3f", "(query, passage)", good, goodMrr / 60);
|
||||
t.line("%-24s hit@1 %d/60 MRR %.3f", "(passage, query)", swapped, swappedMrr / 60);
|
||||
t.line("no error either way: both orders return a number for every pair");
|
||||
}
|
||||
}
|
||||
|
||||
private static int rankAfter(List<String> pids, float[] scores, String gold) {
|
||||
List<Integer> idx = new ArrayList<>();
|
||||
for (int i = 0; i < pids.size(); i++) {
|
||||
idx.add(i);
|
||||
}
|
||||
idx.sort((a, b) -> Float.compare(scores[b], scores[a]));
|
||||
for (int r = 0; r < idx.size(); r++) {
|
||||
if (pids.get(idx.get(r)).equals(gold)) {
|
||||
return r + 1;
|
||||
}
|
||||
}
|
||||
return 0;
|
||||
}
|
||||
|
||||
@Test
|
||||
void scoresAreNotProbabilities() {
|
||||
CrossEncoder ce = Fixture.fp32();
|
||||
float[][] all = Fixture.allScores(ce);
|
||||
try (Transcript t = new Transcript("12-trap-score-scale.txt", "Raw cross-encoder scores for the gold passage and for the best wrong passage")) {
|
||||
List<Float> gold = new ArrayList<>();
|
||||
List<Float> bestWrong = new ArrayList<>();
|
||||
int below0 = 0;
|
||||
for (int qi = 0; qi < KnowledgeBase.QUERIES.size(); qi++) {
|
||||
KnowledgeBase.Query q = KnowledgeBase.QUERIES.get(qi);
|
||||
float g = all[qi][Fixture.pidIndex(q.pid())];
|
||||
float w = -1000;
|
||||
float top = -1000;
|
||||
for (int pi = 0; pi < all[qi].length; pi++) {
|
||||
top = Math.max(top, all[qi][pi]);
|
||||
if (pi != Fixture.pidIndex(q.pid())) {
|
||||
w = Math.max(w, all[qi][pi]);
|
||||
}
|
||||
}
|
||||
gold.add(g);
|
||||
bestWrong.add(w);
|
||||
if (top < 0) {
|
||||
below0++;
|
||||
}
|
||||
}
|
||||
Collections.sort(gold);
|
||||
Collections.sort(bestWrong);
|
||||
t.line("gold passage score : min %.1f, median %.1f, max %.1f", gold.getFirst(), gold.get(30), gold.getLast());
|
||||
t.line("best wrong passage score : min %.1f, median %.1f, max %.1f", bestWrong.getFirst(), bestWrong.get(30), bestWrong.getLast());
|
||||
t.line("questions whose best-scoring passage is below 0: %d of 60", below0);
|
||||
for (float cutoff : new float[] {0f, 2f, 4f}) {
|
||||
long kept = KnowledgeBase.QUERIES.stream().filter(q -> {
|
||||
int qi = KnowledgeBase.QUERIES.indexOf(q);
|
||||
return all[qi][Fixture.pidIndex(q.pid())] >= cutoff;
|
||||
}).count();
|
||||
t.line("a fixed cut-off of %.0f would drop the gold passage for %d of 60 questions", cutoff, 60 - kept);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@Test
|
||||
void truncation() {
|
||||
String answer = "Error E-9999 - Zebra overflow. Reset the zebra counter from the admin console.";
|
||||
String filler = "The audit trail records the user, the time and the address of each change. ".repeat(12);
|
||||
try (Transcript t = new Transcript("13-trap-truncation.txt", "The answer sits at the end of a long passage");
|
||||
CrossEncoder at512 = new CrossEncoder(Models.crossModel(Models.CROSS_FP32), Models.crossTokenizer(), 512, 8, 0);
|
||||
CrossEncoder at64 = new CrossEncoder(Models.crossModel(Models.CROSS_FP32), Models.crossTokenizer(), 64, 8, 0)) {
|
||||
String longPassage = filler + answer;
|
||||
String question = "How do I fix E-9999?";
|
||||
int tokens = at512.encode(question, longPassage).getIds().length;
|
||||
t.line("question + passage = %d tokens", tokens);
|
||||
t.line("score with max length 512 : %.2f", at512.score(question, List.of(longPassage))[0]);
|
||||
t.line("score with max length 64 : %.2f (the answer was cut off before the model saw it)", at64.score(question, List.of(longPassage))[0]);
|
||||
t.line("score of the answer alone : %.2f", at512.score(question, List.of(answer))[0]);
|
||||
long over = KnowledgeBase.PASSAGES.stream().filter(p -> at512.encode("question", p.text()).getIds().length > Fixture.MAX_LEN).count();
|
||||
t.line("passages in this corpus longer than %d tokens: %d of %d", Fixture.MAX_LEN, over, KnowledgeBase.PASSAGES.size());
|
||||
assertThat(at512.score(question, List.of(longPassage))[0]).isGreaterThan(at64.score(question, List.of(longPassage))[0]);
|
||||
}
|
||||
}
|
||||
|
||||
@Test
|
||||
void poolTooSmall() {
|
||||
CrossEncoder ce = Fixture.fp32();
|
||||
try (Transcript t = new Transcript("14-trap-small-pool.txt", "Questions whose gold passage is not in the bi-encoder's top 5")) {
|
||||
int missed = 0, rescued = 0;
|
||||
t.line("%-8s %-62s %-12s %s", "gold", "question", "bi-encoder", "reranked top 50");
|
||||
for (KnowledgeBase.Query q : KnowledgeBase.QUERIES) {
|
||||
int bi = Metrics.rankOf(Fixture.biRanking(q), q.pid());
|
||||
if (bi > 5) {
|
||||
missed++;
|
||||
int after = Metrics.rankOf(Fixture.reranked(ce, q, 50), q.pid());
|
||||
rescued += after <= 5 ? 1 : 0;
|
||||
String text = q.text().length() > 60 ? q.text().substring(0, 59) + "~" : q.text();
|
||||
t.line("%-8s %-62s %-12d %d", q.pid().length() > 8 ? q.pid().substring(0, 8) : q.pid(), text, bi, after);
|
||||
}
|
||||
}
|
||||
t.blank();
|
||||
t.line("%d questions have no answer passage in the top 5; a reranker over the top 5 cannot help them,", missed);
|
||||
t.line("a reranker over the top 50 brings %d of them into the top 5", rescued);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,93 @@
|
||||
package com.ankurm.rerank.support;
|
||||
|
||||
import java.util.ArrayList;
|
||||
import java.util.HashMap;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
|
||||
import com.ankurm.rerank.BiEncoderRetriever;
|
||||
import com.ankurm.rerank.CrossEncoder;
|
||||
import com.ankurm.rerank.KnowledgeBase;
|
||||
import com.ankurm.rerank.Models;
|
||||
import org.springframework.ai.document.Document;
|
||||
import org.springframework.ai.rag.Query;
|
||||
|
||||
/** Models are loaded once per JVM: the embeddings of the 94 passages and the two cross-encoder sessions. */
|
||||
public final class Fixture {
|
||||
|
||||
public static final int MAX_LEN = 256;
|
||||
private static BiEncoderRetriever retriever;
|
||||
private static CrossEncoder fp32;
|
||||
private static CrossEncoder int8;
|
||||
|
||||
private Fixture() {
|
||||
}
|
||||
|
||||
public static synchronized BiEncoderRetriever retriever() {
|
||||
if (retriever == null) {
|
||||
retriever = new BiEncoderRetriever(Models.biEncoder(), KnowledgeBase.PASSAGES, KnowledgeBase.PASSAGES.size());
|
||||
}
|
||||
return retriever;
|
||||
}
|
||||
|
||||
public static synchronized CrossEncoder fp32() {
|
||||
if (fp32 == null) {
|
||||
fp32 = new CrossEncoder(Models.crossModel(Models.CROSS_FP32), Models.crossTokenizer(), MAX_LEN, 32, 0);
|
||||
}
|
||||
return fp32;
|
||||
}
|
||||
|
||||
public static synchronized CrossEncoder int8() {
|
||||
if (int8 == null) {
|
||||
int8 = new CrossEncoder(Models.crossModel(Models.CROSS_INT8), Models.crossTokenizer(), MAX_LEN, 32, 0);
|
||||
}
|
||||
return int8;
|
||||
}
|
||||
|
||||
private static final Map<CrossEncoder, float[][]> ALL_SCORES = new HashMap<>();
|
||||
|
||||
/** Bi-encoder ranking of all 94 passages for a query, best first. */
|
||||
public static List<String> biRanking(KnowledgeBase.Query q) {
|
||||
List<String> out = new ArrayList<>();
|
||||
for (Document d : retriever().retrieve(new Query(q.text()))) {
|
||||
out.add((String) d.getMetadata().get("pid"));
|
||||
}
|
||||
return out;
|
||||
}
|
||||
|
||||
/**
|
||||
* The cross-encoder's score for every (query, passage) pair, indexed [query][passage in KnowledgeBase order].
|
||||
* Reranking the bi-encoder's top N is then "take the top N, sort by these scores", which is exactly what
|
||||
* the post-processor does, without paying for the 94-passage scoring once per pool size.
|
||||
*/
|
||||
public static synchronized float[][] allScores(CrossEncoder ce) {
|
||||
return ALL_SCORES.computeIfAbsent(ce, c -> {
|
||||
List<String> texts = KnowledgeBase.PASSAGES.stream().map(KnowledgeBase.Passage::text).toList();
|
||||
float[][] m = new float[KnowledgeBase.QUERIES.size()][];
|
||||
for (int i = 0; i < m.length; i++) {
|
||||
m[i] = c.score(KnowledgeBase.QUERIES.get(i).text(), texts);
|
||||
}
|
||||
return m;
|
||||
});
|
||||
}
|
||||
|
||||
/** Ranking after cross-encoder reranking of the bi-encoder's top {@code pool}; ranks below the pool keep their order. */
|
||||
public static List<String> reranked(CrossEncoder ce, KnowledgeBase.Query q, int pool) {
|
||||
int qi = KnowledgeBase.QUERIES.indexOf(q);
|
||||
float[] scores = allScores(ce)[qi];
|
||||
List<String> bi = biRanking(q);
|
||||
List<String> head = new ArrayList<>(bi.subList(0, Math.min(pool, bi.size())));
|
||||
head.sort((a, b) -> Float.compare(scores[pidIndex(b)], scores[pidIndex(a)]));
|
||||
head.addAll(bi.subList(Math.min(pool, bi.size()), bi.size()));
|
||||
return head;
|
||||
}
|
||||
|
||||
public static int pidIndex(String pid) {
|
||||
for (int i = 0; i < KnowledgeBase.PASSAGES.size(); i++) {
|
||||
if (KnowledgeBase.PASSAGES.get(i).pid().equals(pid)) {
|
||||
return i;
|
||||
}
|
||||
}
|
||||
throw new IllegalArgumentException(pid);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,50 @@
|
||||
package com.ankurm.rerank.support;
|
||||
|
||||
import java.util.List;
|
||||
import java.util.function.Function;
|
||||
|
||||
import com.ankurm.rerank.KnowledgeBase;
|
||||
|
||||
/** hit@1, hit@3, hit@5 and MRR@10 of a ranker over a set of labelled queries. */
|
||||
public final class Metrics {
|
||||
|
||||
public record Score(int n, int hit1, int hit3, int hit5, double mrr) {
|
||||
|
||||
public String cells() {
|
||||
return String.format("%2d/%-2d %2d/%-2d %2d/%-2d %.3f", hit1, n, hit3, n, hit5, n, mrr);
|
||||
}
|
||||
}
|
||||
|
||||
private Metrics() {
|
||||
}
|
||||
|
||||
/** The 1-based rank of the gold passage in a ranking, or 0 when it is not there. */
|
||||
public static int rankOf(List<String> ranking, String gold) {
|
||||
return ranking.indexOf(gold) + 1;
|
||||
}
|
||||
|
||||
public static Score score(List<KnowledgeBase.Query> queries, Function<KnowledgeBase.Query, List<String>> ranker) {
|
||||
int h1 = 0, h3 = 0, h5 = 0;
|
||||
double mrr = 0;
|
||||
for (KnowledgeBase.Query q : queries) {
|
||||
int rank = rankOf(ranker.apply(q), q.pid());
|
||||
if (rank == 1) {
|
||||
h1++;
|
||||
}
|
||||
if (rank >= 1 && rank <= 3) {
|
||||
h3++;
|
||||
}
|
||||
if (rank >= 1 && rank <= 5) {
|
||||
h5++;
|
||||
}
|
||||
if (rank >= 1 && rank <= 10) {
|
||||
mrr += 1.0 / rank;
|
||||
}
|
||||
}
|
||||
return new Score(queries.size(), h1, h3, h5, mrr / queries.size());
|
||||
}
|
||||
|
||||
public static List<KnowledgeBase.Query> category(String name) {
|
||||
return KnowledgeBase.QUERIES.stream().filter(q -> q.category().equals(name)).toList();
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,26 @@
|
||||
package com.ankurm.rerank.support;
|
||||
|
||||
import java.util.List;
|
||||
import java.util.concurrent.CopyOnWriteArrayList;
|
||||
|
||||
import org.springframework.ai.chat.messages.AssistantMessage;
|
||||
import org.springframework.ai.chat.model.ChatModel;
|
||||
import org.springframework.ai.chat.model.ChatResponse;
|
||||
import org.springframework.ai.chat.model.Generation;
|
||||
import org.springframework.ai.chat.prompt.Prompt;
|
||||
|
||||
/** Records the prompts it is sent and answers "ok". What Spring AI sends is the part this module is responsible for. */
|
||||
public class RecordingChatModel implements ChatModel {
|
||||
|
||||
private final List<Prompt> prompts = new CopyOnWriteArrayList<>();
|
||||
|
||||
public List<Prompt> prompts() {
|
||||
return prompts;
|
||||
}
|
||||
|
||||
@Override
|
||||
public ChatResponse call(Prompt prompt) {
|
||||
prompts.add(prompt);
|
||||
return new ChatResponse(List.of(new Generation(new AssistantMessage("ok"))));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,52 @@
|
||||
package com.ankurm.rerank.support;
|
||||
|
||||
import java.io.IOException;
|
||||
import java.io.PrintWriter;
|
||||
import java.io.StringWriter;
|
||||
import java.nio.file.Files;
|
||||
import java.nio.file.Path;
|
||||
|
||||
/**
|
||||
* Writes a numbered transcript under {@code output/} and echoes it to the console. Every console block
|
||||
* quoted in the article comes out of one of these files verbatim.
|
||||
*/
|
||||
public final class Transcript implements AutoCloseable {
|
||||
|
||||
private final Path path;
|
||||
private final StringWriter buffer = new StringWriter();
|
||||
private final PrintWriter out = new PrintWriter(buffer);
|
||||
|
||||
public Transcript(String fileName, String title) {
|
||||
this.path = Path.of("output", fileName);
|
||||
out.println("# " + title);
|
||||
out.println();
|
||||
}
|
||||
|
||||
public Transcript line(String format, Object... args) {
|
||||
out.println(args.length == 0 ? format : String.format(format, args));
|
||||
return this;
|
||||
}
|
||||
|
||||
public Transcript blank() {
|
||||
out.println();
|
||||
return this;
|
||||
}
|
||||
|
||||
public Transcript section(String heading) {
|
||||
out.println();
|
||||
out.println("--- " + heading + " ---");
|
||||
return this;
|
||||
}
|
||||
|
||||
@Override
|
||||
public void close() {
|
||||
out.flush();
|
||||
try {
|
||||
Files.createDirectories(path.getParent());
|
||||
Files.writeString(path, buffer.toString());
|
||||
} catch (IOException e) {
|
||||
throw new IllegalStateException("could not write " + path, e);
|
||||
}
|
||||
System.out.print(buffer);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user