Generic embedding models don't understand your domain's vocabulary, synonyms, or conceptual relationships — and that gap silently kills retrieval accuracy. This lesson teaches you how to fine-tune a sentence transformer on your own data, from synthetic training pair generation to quantitative evaluation and production deployment.
You've built a solid RAG pipeline. Documents are chunked intelligently, indexed in a vector database, and your LLM produces coherent answers. But something is off — users ask about "cath lab protocols" and the retriever surfaces documents about "cardiac catheterization procedures" instead of the exact nursing checklists they need. Or a legal team queries for "force majeure clauses" and the system returns broadly related contract language instead of the specific indemnification sections that matter. The semantic gap isn't a retrieval strategy problem. It's a representation problem — the embedding model doesn't understand your domain the way your users do.
Generic embedding models like text-embedding-3-small or all-MiniLM-L6-v2 are trained on massive, diverse internet corpora. They're remarkably capable across broad tasks, but they've never seen your internal clinical documentation, your proprietary legal contracts, or your specialized manufacturing runbooks. The vocabulary your domain uses, the relationships between concepts, and the way your users phrase queries — these are invisible to a general-purpose model. Fine-tuning an embedding model on your domain data teaches it the specific semantic relationships that matter for your use case, often improving retrieval precision by 20–40% on specialized corpora.
By the end of this lesson, you'll have a complete, production-grade workflow for fine-tuning a sentence transformer embedding model on domain-specific data and deploying it in a RAG pipeline.
What you'll learn:
sentence-transformers libraryThis lesson assumes you're comfortable with:
You'll need: Python 3.10+, a CUDA-capable GPU (or Google Colab with a T4/A100), and the packages listed at the start of each section.
Before writing a single line of training code, it's worth understanding why this problem exists — because it shapes every decision you'll make.
Embedding models are trained with a contrastive objective: semantically similar text should produce vectors that are close together in high-dimensional space; dissimilar text should be pushed apart. The model learns "similarity" from its training data. When you use all-mpnet-base-v2 on a biomedical corpus, "similar" still means what Wikipedia and Common Crawl taught it — which may not align with how cardiologists or oncologists define conceptual similarity.
Consider a concrete failure case. Suppose you're building a RAG system over internal policy documents for a financial services firm. A user asks:
"What's our exposure on naked short positions under the new T+1 rules?"
A general embedding model maps "naked short positions" into a neighborhood populated with generic finance articles about short selling. But your corpus contains a specific risk management policy titled "Intraday Margin Requirements for Uncovered Short Sales" — which is the exact document this user needs. The term "naked" vs "uncovered" is a domain synonym; "T+1" refers to SEC settlement rule changes from 2024. A model that's never learned these equivalences in your vocabulary space will miss the match.
The fundamental issue is that the model's learned representation of the embedding space doesn't reflect the semantic structure of your specific domain. Fine-tuning reshapes that structure to fit your data.
Key insight: You're not training the model to memorize your documents. You're teaching it which concepts are semantically equivalent in your domain. The retrieval improvement comes from a better-calibrated embedding space, not from the model "knowing" the answers.
There's also a subtler issue: query-document asymmetry. Users phrase queries conversationally; your documents are dense, formal text. A model that's seen millions of natural language queries paired with relevant passages handles this bridge well — for general topics. But your users may write technical queries that look nothing like how your documents are written, and the bridge needs to be learned domain-specifically.
Fine-tuning for retrieval is different from fine-tuning an LLM for generation. You're optimizing the geometry of embedding space, not token prediction probabilities. The standard approach uses contrastive loss — specifically, MultipleNegativesRankingLoss (MNRL), which is the workhorse of modern embedding model training.
Here's the intuition: You have a batch of (query, positive_document) pairs. For each query in the batch, the corresponding positive document is the target. All other documents in the same batch serve as in-batch negatives — documents that are not relevant to this query but look structurally similar (they're also real domain documents). The model is trained to rank the positive document above all negatives.
Mathematically, MNRL optimizes:
loss = -log( exp(sim(q, d+) / τ) / Σ exp(sim(q, dᵢ) / τ) )
Where q is the query embedding, d+ is the positive document embedding, dᵢ are all documents in the batch, and τ is a temperature parameter. This is a cross-entropy loss over the similarity scores.
The critical implication: batch size matters enormously. With a batch size of 64, each query sees 63 in-batch negatives. With 256, it sees 255. Larger batches provide harder, more informative negatives and generally produce better models — up to the limit of your GPU memory. This is why fine-tuning embedding models benefits significantly from larger GPUs.
Tip: If you're GPU-memory constrained, the
sentence-transformerslibrary supports gradient caching and GradCache, which lets you compute MNRL with effectively larger batches by accumulating gradients across mini-batches. We'll show this in the training code below.
Beyond MNRL, you can also use TripletLoss (query, positive, hard negative) or CosineSimilarityLoss (pairs with similarity scores). For retrieval, MNRL is almost always the right starting point because it requires only positive pairs — which are much easier to collect than explicit hard negatives.
This is where most practitioners either get stuck or make expensive mistakes. The quality of your training data is the single biggest factor in fine-tuning success. Let's walk through three progressively better approaches.
If your organization has existing search logs — users queried something, clicked on a document, and we know which — you have gold labels. Parse these into (query, document_text) pairs.
import pandas as pd
from pathlib import Path
# Hypothetical: a CSV of search logs with query, doc_id, clicked (0/1)
logs = pd.read_csv("search_logs.csv")
positive_pairs = logs[logs["clicked"] == 1][["query", "doc_id"]]
# Join to your document corpus
corpus = pd.read_parquet("document_corpus.parquet") # doc_id, text columns
training_data = positive_pairs.merge(corpus, on="doc_id")[["query", "text"]]
print(f"Training pairs: {len(training_data):,}")
# e.g., Training pairs: 8,432
Real logs are noisy — clicks aren't perfect relevance signals. Filter out sessions where the user clicked and immediately bounced, or where the click was the only result (no ranking signal). Even 5,000–10,000 clean pairs from logs will meaningfully improve your model.
When you don't have search logs, you can use an LLM to generate synthetic queries for each document chunk. This technique — sometimes called GPL (Generative Pseudo Labeling) — is surprisingly effective.
from openai import OpenAI
import json
from tqdm import tqdm
client = OpenAI()
def generate_queries_for_chunk(chunk_text: str, n_queries: int = 3) -> list[str]:
"""Generate realistic user queries that this chunk would answer."""
prompt = f"""You are generating training data for a retrieval system used by healthcare professionals.
Given the following clinical document excerpt, generate {n_queries} realistic questions that a nurse, physician, or clinical administrator might ask that this passage would directly answer.
Requirements:
- Use natural, conversational phrasing a real user would type
- Include domain-specific terminology as appropriate
- Vary the phrasing across questions (different angles on the same content)
- Do NOT include questions the passage cannot answer
Document excerpt:
{chunk_text}
Return a JSON array of {n_queries} question strings. No other text."""
response = client.chat.completions.create(
model="gpt-4o-mini", # Cost-efficient for bulk generation
messages=[{"role": "user", "content": prompt}],
response_format={"type": "json_object"},
temperature=0.7
)
try:
result = json.loads(response.choices[0].message.content)
# Handle various JSON structures the model might return
if isinstance(result, list):
return result
elif isinstance(result, dict):
# Try common keys
for key in ["questions", "queries", "items"]:
if key in result and isinstance(result[key], list):
return result[key]
except (json.JSONDecodeError, KeyError):
pass
return []
# Load your chunked documents
chunks = []
with open("clinical_chunks.jsonl") as f:
for line in f:
chunks.append(json.loads(line))
# Generate training pairs
training_pairs = []
for chunk in tqdm(chunks[:2000]): # Start with a manageable subset
queries = generate_queries_for_chunk(chunk["text"], n_queries=3)
for q in queries:
if q and len(q.strip()) > 10:
training_pairs.append({
"query": q.strip(),
"positive": chunk["text"],
"doc_id": chunk["id"]
})
print(f"Generated {len(training_pairs):,} training pairs")
# Generated 5,847 training pairs (some chunks produce fewer than 3)
# Save
with open("synthetic_training_pairs.jsonl", "w") as f:
for pair in training_pairs:
f.write(json.dumps(pair) + "\n")
Warning: Synthetic pairs have a systematic bias: the LLM generates questions that are obviously answerable from the chunk. This makes training pairs somewhat easier than real user queries, which are often ambiguous or use different vocabulary. Supplement with at least some real queries when possible.
In-batch negatives from MNRL are random negatives — other documents in the batch. For harder training, you want hard negatives: documents that are topically similar but not the right answer. These make the model learn finer distinctions.
from sentence_transformers import SentenceTransformer
import numpy as np
import faiss
# Use a baseline model to find hard negatives
baseline_model = SentenceTransformer("all-MiniLM-L6-v2")
# Embed all document chunks
print("Embedding corpus...")
corpus_texts = [chunk["text"] for chunk in chunks]
corpus_embeddings = baseline_model.encode(
corpus_texts,
batch_size=128,
show_progress_bar=True,
normalize_embeddings=True
)
# Build a FAISS index for fast retrieval
dimension = corpus_embeddings.shape[1]
index = faiss.IndexFlatIP(dimension) # Inner product = cosine similarity for normalized vectors
index.add(corpus_embeddings.astype(np.float32))
# For each training pair, find hard negatives
def find_hard_negatives(query_text: str, positive_doc_id: str, k: int = 5) -> list[str]:
"""Find documents that are similar to the query but aren't the positive."""
query_embedding = baseline_model.encode(
[query_text], normalize_embeddings=True
)
scores, indices = index.search(query_embedding.astype(np.float32), k + 5)
hard_negatives = []
for idx in indices[0]:
if chunks[idx]["id"] != positive_doc_id:
hard_negatives.append(chunks[idx]["text"])
if len(hard_negatives) >= k:
break
return hard_negatives
# Add hard negatives to training data
enriched_pairs = []
for pair in training_pairs[:500]: # Sample for demonstration
hard_negs = find_hard_negatives(pair["query"], pair["doc_id"], k=3)
enriched_pairs.append({
**pair,
"hard_negatives": hard_negs
})
Hard negatives are used with TripletLoss or with MNRL's extended variant that accepts explicit negatives. For most production use cases, starting with MNRL and in-batch negatives is the right call — add hard negatives if your first fine-tuned model still struggles with similar-but-not-quite-right retrievals.
Now we get to the actual training. We'll use the sentence-transformers library, which wraps Hugging Face transformers with the contrastive training infrastructure you need.
pip install sentence-transformers datasets accelerate
from sentence_transformers import SentenceTransformer, InputExample, losses
from sentence_transformers.evaluation import InformationRetrievalEvaluator
from torch.utils.data import DataLoader
import json
import random
from pathlib import Path
# --- Configuration ---
BASE_MODEL = "sentence-transformers/all-mpnet-base-v2" # Strong baseline
OUTPUT_DIR = "./clinical-embedding-v1"
BATCH_SIZE = 64 # Increase if VRAM allows; more = harder in-batch negatives
NUM_EPOCHS = 3
WARMUP_RATIO = 0.1
LEARNING_RATE = 2e-5
# --- Load Training Data ---
training_pairs = []
with open("synthetic_training_pairs.jsonl") as f:
for line in f:
training_pairs.append(json.loads(line))
# Shuffle for good measure
random.shuffle(training_pairs)
# Split: 90% train, 10% validation
split = int(len(training_pairs) * 0.9)
train_pairs = training_pairs[:split]
val_pairs = training_pairs[split:]
print(f"Train: {len(train_pairs):,} | Val: {len(val_pairs):,}")
# --- Convert to InputExample format ---
train_examples = [
InputExample(texts=[pair["query"], pair["positive"]])
for pair in train_pairs
]
# --- Load Model ---
model = SentenceTransformer(BASE_MODEL)
# --- Data Loader ---
train_dataloader = DataLoader(
train_examples,
shuffle=True,
batch_size=BATCH_SIZE,
drop_last=True # MNRL benefits from consistent batch sizes
)
# --- Loss Function ---
train_loss = losses.MultipleNegativesRankingLoss(model=model)
# --- Validation Evaluator ---
# Build evaluation corpus from val pairs
val_queries = {str(i): pair["query"] for i, pair in enumerate(val_pairs)}
val_corpus = {str(i): pair["positive"] for i, pair in enumerate(val_pairs)}
val_relevant_docs = {str(i): {str(i)} for i in range(len(val_pairs))}
evaluator = InformationRetrievalEvaluator(
queries=val_queries,
corpus=val_corpus,
relevant_docs=val_relevant_docs,
name="clinical-val",
score_functions={"cos_sim": lambda x, y: (x @ y.T)}
)
# --- Warmup Steps ---
total_steps = len(train_dataloader) * NUM_EPOCHS
warmup_steps = int(total_steps * WARMUP_RATIO)
print(f"Total steps: {total_steps:,} | Warmup steps: {warmup_steps:,}")
# --- Train ---
model.fit(
train_objectives=[(train_dataloader, train_loss)],
evaluator=evaluator,
epochs=NUM_EPOCHS,
warmup_steps=warmup_steps,
optimizer_params={"lr": LEARNING_RATE},
output_path=OUTPUT_DIR,
save_best_model=True,
evaluation_steps=500,
show_progress_bar=True,
use_amp=True # Mixed precision — significant speedup on modern GPUs
)
print(f"Fine-tuning complete. Best model saved to {OUTPUT_DIR}")
Note: The
InformationRetrievalEvaluatorcomputes NDCG@10, MRR@10, and MAP across your validation set. Watch these metrics during training — if NDCG@10 stops improving after epoch 1 but continues in epoch 2, that's expected. If it drops, you may be overfitting. Thesave_best_model=Trueflag ensures you get the checkpoint from peak validation performance.
Your base model choice matters. Here's a practical guide:
| Base Model | Size | Speed | Notes |
|---|---|---|---|
all-MiniLM-L6-v2 |
80MB | Very fast | Good starting point; may plateau faster |
all-mpnet-base-v2 |
420MB | Moderate | Strong baseline, excellent for most domains |
BAAI/bge-base-en-v1.5 |
440MB | Moderate | State-of-the-art retrieval out of the box; great to fine-tune |
BAAI/bge-large-en-v1.5 |
1.3GB | Slow | Maximum quality; fine-tune if you have budget |
intfloat/e5-base-v2 |
440MB | Moderate | Strong cross-lingual capability if needed |
For most production fine-tuning, start with BAAI/bge-base-en-v1.5 — it's already SOTA for retrieval and fine-tunes well. The smaller MiniLM models are tempting for speed but have a lower ceiling.
Tip: If you're operating in a regulated environment with data sovereignty requirements, check that your base model's license permits commercial use and that you can host it entirely on-premises. Both
bgeandmpnetmodels use Apache 2.0 licenses — safe for commercial deployment.
Validation loss during training tells you the model is learning. But what you really care about is downstream retrieval quality — does it actually retrieve better documents for real queries? This requires a more rigorous evaluation setup.
from sentence_transformers import SentenceTransformer
import numpy as np
import faiss
from tqdm import tqdm
import json
def build_retrieval_index(model: SentenceTransformer, corpus: list[dict]) -> tuple:
"""Embed corpus and return FAISS index + metadata."""
texts = [doc["text"] for doc in corpus]
embeddings = model.encode(
texts,
batch_size=128,
show_progress_bar=True,
normalize_embeddings=True
)
dim = embeddings.shape[1]
index = faiss.IndexFlatIP(dim)
index.add(embeddings.astype(np.float32))
return index, embeddings
def evaluate_retrieval(
model: SentenceTransformer,
corpus: list[dict],
eval_set: list[dict], # [{"query": ..., "relevant_doc_ids": [...]}]
k_values: list[int] = [1, 5, 10, 20]
) -> dict:
"""Compute Recall@k and NDCG@k for a retrieval model."""
index, _ = build_retrieval_index(model, corpus)
results = {f"recall@{k}": [] for k in k_values}
results.update({f"ndcg@{k}": [] for k in k_values})
for item in tqdm(eval_set, desc="Evaluating"):
query_embedding = model.encode(
[item["query"]],
normalize_embeddings=True
).astype(np.float32)
max_k = max(k_values)
scores, indices = index.search(query_embedding, max_k)
retrieved_ids = [corpus[idx]["id"] for idx in indices[0]]
relevant_ids = set(item["relevant_doc_ids"])
for k in k_values:
top_k = retrieved_ids[:k]
hits = sum(1 for doc_id in top_k if doc_id in relevant_ids)
# Recall@k
recall = hits / len(relevant_ids) if relevant_ids else 0
results[f"recall@{k}"].append(recall)
# NDCG@k
dcg = sum(
1.0 / np.log2(rank + 2)
for rank, doc_id in enumerate(top_k)
if doc_id in relevant_ids
)
ideal_dcg = sum(
1.0 / np.log2(rank + 2)
for rank in range(min(len(relevant_ids), k))
)
ndcg = dcg / ideal_dcg if ideal_dcg > 0 else 0
results[f"ndcg@{k}"].append(ndcg)
return {metric: np.mean(values) for metric, values in results.items()}
# Load models
baseline_model = SentenceTransformer("BAAI/bge-base-en-v1.5")
finetuned_model = SentenceTransformer("./clinical-embedding-v1")
# Load held-out test corpus and evaluation queries
# These should be COMPLETELY SEPARATE from training data
with open("test_corpus.jsonl") as f:
test_corpus = [json.loads(line) for line in f]
with open("test_eval_set.jsonl") as f:
eval_set = [json.loads(line) for line in f]
print(f"Test corpus: {len(test_corpus):,} documents | Eval queries: {len(eval_set):,}")
# Evaluate both
print("\n--- Baseline Model ---")
baseline_results = evaluate_retrieval(baseline_model, test_corpus, eval_set)
for metric, value in sorted(baseline_results.items()):
print(f" {metric}: {value:.4f}")
print("\n--- Fine-Tuned Model ---")
finetuned_results = evaluate_retrieval(finetuned_model, test_corpus, eval_set)
for metric, value in sorted(finetuned_results.items()):
print(f" {metric}: {value:.4f}")
# Delta
print("\n--- Improvement ---")
for metric in baseline_results:
delta = finetuned_results[metric] - baseline_results[metric]
pct = (delta / baseline_results[metric]) * 100 if baseline_results[metric] > 0 else 0
print(f" {metric}: {delta:+.4f} ({pct:+.1f}%)")
A realistic output for a healthcare domain fine-tune might look like:
--- Improvement ---
ndcg@10: +0.0832 (+18.6%)
recall@5: +0.1124 (+23.4%)
recall@10: +0.0891 (+15.2%)
recall@20: +0.0643 (+9.8%)
These are meaningful, real-world gains. Recall@5 improving by 23% means that in roughly 1 out of 4 queries where the baseline failed, your fine-tuned model now retrieves the right document — often before your reranker or cross-encoder reranking step even has a chance to see it.
Warning: Don't evaluate on data that overlaps with your training set. This is the #1 evaluation mistake in embedding fine-tuning. If you generated synthetic queries from your entire corpus and then evaluate on that same corpus, you're measuring memorization, not generalization. Always use a held-out document set for evaluation.
Swapping the embedding model into an existing pipeline is straightforward, but there's a critical operational detail: you must re-embed your entire document corpus with the new model. Embeddings from different models are not comparable — you can't mix base model vectors with fine-tuned model vectors in the same index.
from sentence_transformers import SentenceTransformer
import numpy as np
import json
# Load your fine-tuned model
embedding_model = SentenceTransformer("./clinical-embedding-v1")
# If you need it for inference in a different service or want to export
# the model to Hugging Face Hub for team use:
# embedding_model.save_to_hub("your-org/clinical-embedding-v1", private=True)
class DomainEmbedder:
"""Drop-in replacement embedder for your RAG pipeline."""
def __init__(self, model_path: str, batch_size: int = 64):
self.model = SentenceTransformer(model_path)
self.batch_size = batch_size
self.dimension = self.model.get_sentence_embedding_dimension()
def embed_documents(self, texts: list[str]) -> np.ndarray:
"""Embed a list of document chunks for indexing."""
return self.model.encode(
texts,
batch_size=self.batch_size,
normalize_embeddings=True,
show_progress_bar=len(texts) > 100
)
def embed_query(self, query: str) -> np.ndarray:
"""Embed a single user query for retrieval."""
# Some models (e5, bge) benefit from query prefix
# bge-base recommends prepending "Represent this sentence for searching relevant passages: "
prefixed = f"Represent this sentence for searching relevant passages: {query}"
return self.model.encode(
[prefixed],
normalize_embeddings=True
)[0]
def get_dimension(self) -> int:
return self.dimension
# Re-index your corpus
embedder = DomainEmbedder("./clinical-embedding-v1")
# Load chunks (same chunks as before — only the embeddings change)
with open("clinical_chunks.jsonl") as f:
chunks = [json.loads(line) for line in f]
texts = [chunk["text"] for chunk in chunks]
print(f"Re-embedding {len(texts):,} chunks with dimension {embedder.get_dimension()}...")
embeddings = embedder.embed_documents(texts)
print(f"Embeddings shape: {embeddings.shape}") # e.g., (15000, 768)
# Save embeddings (format depends on your vector store)
np.save("clinical_embeddings_v1.npy", embeddings)
print("Done. Update your vector store with these new embeddings.")
For production pipelines using a managed vector database like Pinecone or Weaviate, you'd upsert these new vectors with the same IDs as your existing documents. If you're using pgvector, it's a straightforward UPDATE or COPY operation. The vector database implementation guide covers the specifics for each platform.
Note: Re-indexing a large corpus takes time. For a corpus of 100K chunks, expect 10–30 minutes on a GPU and 1–3 hours on CPU. Plan for a brief indexing window when deploying. For live production systems, consider a blue-green deployment: build the new index while the old one serves traffic, then swap atomically.
The query-time integration is even simpler:
# Before (using a generic model via LangChain or similar)
from langchain.embeddings import HuggingFaceEmbeddings
old_embeddings = HuggingFaceEmbeddings(model_name="BAAI/bge-base-en-v1.5")
# After (your fine-tuned model — same interface)
new_embeddings = HuggingFaceEmbeddings(
model_name="./clinical-embedding-v1",
encode_kwargs={"normalize_embeddings": True}
)
# Everything else in your pipeline stays identical
Let's do a complete, end-to-end exercise using a real public dataset. The CUAD (Contract Understanding Atticus Dataset) contains commercial contracts with expert-annotated clauses — a perfect proxy for the kind of specialized legal corpus many enterprises work with.
# Install dependencies
# pip install sentence-transformers datasets faiss-cpu openai tqdm
from datasets import load_dataset
from sentence_transformers import SentenceTransformer, InputExample, losses
from sentence_transformers.evaluation import InformationRetrievalEvaluator
from torch.utils.data import DataLoader
import random
import json
from openai import OpenAI
client = OpenAI()
# --- Step 1: Load CUAD and prepare document chunks ---
print("Loading CUAD dataset...")
cuad = load_dataset("theatticusproject/cuad-qa", split="train")
# Each example has: context (contract text), question (clause query), answers
# We'll use questions as queries and contexts as documents
examples = list(cuad)
random.shuffle(examples)
# Use first 3000 examples for training, last 300 for evaluation
train_examples_raw = examples[:3000]
eval_examples_raw = examples[3000:3300]
print(f"Train: {len(train_examples_raw)} | Eval: {len(eval_examples_raw)}")
# --- Step 2: Build training pairs ---
# CUAD questions are already expert-written — perfect training queries!
# No need for LLM generation here since we have real labeled data
train_input_examples = []
for ex in train_examples_raw:
if ex["answers"]["text"]: # Only use examples with definitive answers
train_input_examples.append(
InputExample(
texts=[ex["question"], ex["context"][:512]] # Truncate context
)
)
print(f"Training pairs: {len(train_input_examples):,}")
# --- Step 3: Build evaluation set ---
eval_queries = {}
eval_corpus = {}
eval_relevant = {}
for i, ex in enumerate(eval_examples_raw):
if ex["answers"]["text"]:
q_id = f"q_{i}"
d_id = f"d_{i}"
eval_queries[q_id] = ex["question"]
eval_corpus[d_id] = ex["context"][:512]
eval_relevant[q_id] = {d_id}
print(f"Eval queries: {len(eval_queries):,}")
# --- Step 4: Fine-tune ---
BASE_MODEL = "BAAI/bge-base-en-v1.5"
model = SentenceTransformer(BASE_MODEL)
train_dataloader = DataLoader(
train_input_examples,
shuffle=True,
batch_size=32,
drop_last=True
)
train_loss = losses.MultipleNegativesRankingLoss(model=model)
evaluator = InformationRetrievalEvaluator(
queries=eval_queries,
corpus=eval_corpus,
relevant_docs=eval_relevant,
name="cuad-eval"
)
total_steps = len(train_dataloader) * 2 # 2 epochs
warmup_steps = total_steps // 10
model.fit(
train_objectives=[(train_dataloader, train_loss)],
evaluator=evaluator,
epochs=2,
warmup_steps=warmup_steps,
optimizer_params={"lr": 2e-5},
output_path="./legal-embedding-cuad",
save_best_model=True,
evaluation_steps=200,
use_amp=True,
show_progress_bar=True
)
print("Done! Check ./legal-embedding-cuad for your fine-tuned model.")
After training, run the evaluation harness from the previous section comparing BAAI/bge-base-en-v1.5 against ./legal-embedding-cuad on a held-out CUAD test set. You should see NDCG@10 improve by roughly 12–25%, depending on your GPU, batch size, and training duration.
This fine-tuned model can drop directly into a legal RAG pipeline. If you're exploring more sophisticated retrieval strategies — like dynamically adjusting how many chunks to retrieve based on query complexity — combining a domain-tuned embedding model with adaptive retrieval techniques compounds the benefit.
Symptom: Your fine-tuned model shows massive improvement on validation (50%+ NDCG gain) but zero improvement in production.
Cause: Queries in your evaluation set were generated from the same documents as your training set. The model learned to overfit the synthetic query patterns.
Fix: Always hold out a document partition before generating training pairs. Queries for evaluation should come from documents the model never trained on.
Symptom: Model improves on domain retrieval but degrades on general queries your pipeline also handles.
Cause: With a high learning rate and many epochs, the model "forgets" its general representations while learning domain-specific ones.
Fix: Reduce learning rate to 1e-5 or lower. Use a warmup ratio of at least 10%. Consider adding a small percentage of general-purpose MSMARCO training pairs to your dataset to anchor general semantics. This is the embedding fine-tuning equivalent of regularization.
Symptom: bge or e5 models perform worse after fine-tuning than expected, or worse than the baseline on some queries.
Cause: bge and e5 models are designed with asymmetric query/document encoding — queries get a prefix ("Represent this sentence for searching...") while documents don't. If you forget the prefix at inference time, your vectors are misaligned.
Fix: Always wrap your fine-tuned model's query encoding with the appropriate prefix. Build it into your embed_query() method as shown in the integration code above so it's never forgotten.
Symptom: Fine-tuning on a small GPU with batch_size=8 produces disappointing results — barely better than the baseline.
Cause: MNRL with tiny batches has almost no useful in-batch negatives. With 8 samples per batch, each query sees only 7 negatives — most of which are random and easy to distinguish from the positive. The model doesn't learn hard discriminations.
Fix: Use gradient caching to simulate larger batches:
from sentence_transformers.losses import CachedMultipleNegativesRankingLoss
# Drop-in replacement for MultipleNegativesRankingLoss
# mini_batch_size is the actual GPU batch; the effective batch is larger
train_loss = CachedMultipleNegativesRankingLoss(
model=model,
mini_batch_size=16 # Real GPU batch; accumulates gradients for effective batch of 64+
)
Fine-tuning your embedding model significantly improves first-stage retrieval, but it doesn't replace the value of a cross-encoder reranker. Think of them as complementary layers: the embedding model retrieves a good candidate set from a massive corpus quickly; the reranker precisely orders the top-K candidates. After fine-tuning your embedding model, revisit your reranking configuration — you may find you can retrieve fewer initial candidates and still maintain the same final answer quality, which reduces latency.
You've now seen the complete arc of domain-specific embedding fine-tuning: understanding why generic models fail, building training data from search logs or LLM-generated synthetic pairs, optimizing contrastive loss with the right batch sizes, evaluating rigorously on held-out data, and deploying with a clean re-indexing strategy.
The key ideas to carry forward:
Where to go next:
A domain-tuned embedding model is one of the highest-leverage improvements you can make to a RAG system. It's not a one-time investment either — as your corpus grows and evolves, periodic re-fine-tuning on fresh data keeps your retrieval quality sharp. Build the training pipeline once, and running a new fine-tuning job becomes routine maintenance rather than a research project.