Most RAG projects that "work in the demo and fail in the pilot" do not have a generation problem. They have a retrieval problem: the top-k chunks handed to the model simply do not contain the answer. Swapping in a bigger LLM cannot fix that, and neither can more prompt engineering.
The two highest-leverage fixes are domain-adapting the embedding model that builds your index, and adding a cross-encoder reranker in front of the generator. This tutorial walks through both on SageMaker AI: building training pairs from your own corpus, fine-tuning a bi-encoder with a SageMaker training job, deploying a reranker endpoint, and measuring whether any of it actually helped.
It assumes you already have a retrieval pipeline of the kind described in Production RAG on AWS. Everything uses the SageMaker Python SDK (pip install "sagemaker>=2.200").
First: measure retrieval on its own
Do not evaluate retrieval by reading generated answers. Build a small labelled set — 150 to 300 real questions, each mapped to the chunk IDs that genuinely answer it — and score the retriever alone.
Three numbers are enough:
- Recall@k — fraction of questions where at least one correct chunk appears in the top k. This is the ceiling on your whole system.
- MRR@10 — how high the first correct chunk ranks.
- nDCG@10 — rank quality when several chunks are relevant.
def recall_at_k(results, gold, k=10):
hits = 0
for qid, ranked in results.items():
if set(ranked[:k]) & set(gold[qid]):
hits += 1
return hits / len(results)
Where do the questions come from? Support tickets, search logs and sales emails are better than anything synthetic. If you have none, generate candidate questions from your chunks with an LLM, then have a human delete the ones that are not realistic. The deletion step is what makes the set usable.
Record the baseline with a stock embedding model before you change anything. Teams routinely find Recall@10 around 0.6 on technical corpora, which means four questions in ten are unanswerable no matter how good the generator is.
Mining training pairs from your corpus
A bi-encoder is fine-tuned on (query, positive passage) pairs, ideally with hard negatives — passages that look similar but do not answer the query. Random negatives teach the model almost nothing.
A pragmatic pipeline:
- For each chunk, prompt an LLM to write two or three questions answerable only from that chunk. Bedrock or a self-hosted endpoint both work; run it as a SageMaker batch transform or Processing job rather than a laptop loop.
- Filter: discard questions that are answerable from a generic chunk, and discard chunks that produce near-duplicate questions.
- For each query, retrieve the top 30 chunks with your current embedding model, drop the known positive, and keep ranks 5–25 as hard negatives. The very top hits are often unlabelled positives, so skipping them avoids training against the truth.
Store the result as JSONL in S3:
{"query": "What is the retention window for audit logs?", "positive": "Audit events are retained for 400 days ...", "negatives": ["Application logs stream to CloudWatch ...", "..."]}
Hold out 10–20% by document, not by row, or near-duplicate chunks leak across the split and your metrics will lie.
Fine-tuning the bi-encoder
sentence-transformers v3+ has a trainer that handles multiple-negatives-ranking loss cleanly. Write a train.py that reads the JSONL, trains, and saves to /opt/ml/model:
from sentence_transformers import SentenceTransformer, SentenceTransformerTrainer, losses
from datasets import load_dataset
model = SentenceTransformer("BAAI/bge-base-en-v1.5")
ds = load_dataset("json", data_files={"train": "/opt/ml/input/data/train/train.jsonl"})["train"]
loss = losses.MultipleNegativesRankingLoss(model)
trainer = SentenceTransformerTrainer(model=model, train_dataset=ds, loss=loss)
trainer.train()
model.save("/opt/ml/model")
Then launch it:
from sagemaker.huggingface import HuggingFace
est = HuggingFace(
entry_point="train.py",
source_dir="src",
role=role,
instance_type="ml.g5.2xlarge",
instance_count=1,
transformers_version="4.49",
pytorch_version="2.5",
py_version="py311",
hyperparameters={"epochs": 2, "batch_size": 64, "lr": 2e-5},
use_spot_instances=True,
max_wait=7200,
)
est.fit({"train": "s3://your-bucket/rag-pairs/"})
Practical notes:
- With MultipleNegativesRankingLoss, batch size is a hyperparameter that matters more than learning rate — other rows in the batch act as in-batch negatives. Use the largest batch that fits; enable gradient checkpointing before you shrink it.
- Two epochs is usually plenty. Embedding models overfit fast and lose general-language ability — a failure mode that only shows up on queries phrased unlike your training set.
- A single g5 run on 20k pairs takes well under an hour, and on managed spot it costs a few dollars. There is no reason to skip this experiment.
- Keep the base model pinned and version the output in the Model Registry so you can tie an index back to the exact encoder that built it.
Re-embedding is mandatory. A fine-tuned encoder is not compatible with vectors produced by the old one. Build a new index (or a new OpenSearch collection), backfill with a batch transform job, and cut over behind a flag so you can roll back.
Adding a cross-encoder reranker
A bi-encoder embeds query and passage separately, which is what makes ANN search fast. A cross-encoder scores the pair jointly and is far more accurate — and far too slow to run over a whole corpus. So you use both: retrieve 50 candidates with the bi-encoder, rerank them to a top 5, and send only those to the LLM.
Deploy a reranker as its own small endpoint:
from sagemaker.huggingface import HuggingFaceModel
reranker = HuggingFaceModel(
model_data="s3://your-bucket/reranker/model.tar.gz",
role=role,
transformers_version="4.49",
pytorch_version="2.5",
py_version="py311",
)
predictor = reranker.deploy(
initial_instance_count=1,
instance_type="ml.g5.xlarge",
endpoint_name="rag-reranker",
)
BAAI/bge-reranker-v2-m3 is a reasonable starting point, and it can be fine-tuned on the same mined pairs with a binary cross-entropy objective. Budget roughly 40–120 ms to rerank 50 passages on a g5.xlarge; keep the candidate count in config so you can trade latency against quality without redeploying.
If the reranker endpoint is idle much of the day, serverless inference or an inference component with scale-to-zero keeps it from billing around the clock — see Inference Components and Scale-to-Zero.
Hybrid search before you fine-tune anything
One caveat worth stating plainly: on corpora full of part numbers, error codes and proper nouns, BM25 keyword search plus reciprocal rank fusion often beats a fine-tuned dense retriever, and costs a day of work instead of a week. OpenSearch supports both in one query. Run that experiment first, then fine-tune on top of the hybrid pipeline rather than instead of it.
Prove it moved
Re-run the same labelled set through four configurations and put the numbers in a table:
| Config | Recall@10 | nDCG@10 | p95 latency |
|---|---|---|---|
| Stock embeddings | baseline | baseline | — |
| + hybrid BM25 fusion | |||
| Fine-tuned embeddings | |||
| Fine-tuned + reranker |
Then check the end-to-end answer quality with an LLM-as-judge harness, because better retrieval occasionally surfaces contradictory passages the generator handles badly. Our guide to automated LLM evaluation gates with fmeval covers wiring that into CI so a retriever change cannot ship without a scorecard.
Finally, log query, retrieved IDs, reranked IDs and thumbs-up/down from production. That log is next quarter's training set, and it is the asset that compounds — the model checkpoint is not.
Where this usually lands
For most teams the order of operations is: measure retrieval in isolation, add hybrid search, add a reranker, and only then fine-tune the embedding model. The first three are days of work; the fourth is the one that needs a labelled set you trust.
Stuck at a Recall@10 that will not move, or deciding whether to fine-tune or just rerank? Talk to us — retrieval quality is a short, well-bounded engagement.