Posted inAI

Fine-Tuning Domain Embedding Models with Sentence-Transformers for Vietnamese RAG

AI tutorial - IT technology blog
AI tutorial - IT technology blog

Context & Why

Six months ago, our team launched an internal RAG assistant to navigate over 45,000 pages of Vietnamese tax laws and corporate compliance circulars. We kicked off with popular multilingual baselines like intfloat/multilingual-e5-base and BAAI/bge-m3. Both fared well on standard generic benchmarks. However, the moment legal analysts pushed queries packed with local administrative phrasing, regional tax codes, and compound terms, retrieval fell apart.

General-purpose bi-encoders build dense vectors from broad web scrapes. When users ran searches like “quyết toán thuế thu nhập cá nhân đối với chuyên gia nước ngoài”, off-the-shelf models matched superficial keywords rather than legal intent. Junk context flooded the prompt window, triggering subtle hallucinations in downstream GPT-4 and Claude summaries.

We faced a tough tradeoff. We could bolt a cross-encoder reranker onto every request, or fine-tune the bi-encoder directly on in-domain pairs. Rerankers wrecked our latency budget, pushing p99 response times past 410ms on our 8-core CPU instances. In contrast, fine-tuning the base bi-encoder with sentence-transformers kept latency down to 22ms while driving Top-5 hit rate from 61.4% to 89.2%.

Fine-tuning is practical engineering, not just academic theory. It gives you enterprise-grade accuracy in specialized languages without blowing up your cloud infrastructure bill.

Installation

Starting with version 3, sentence-transformers integrates natively with Hugging Face’s Trainer. This overhaul brings built-in support for distributed runs, FP16/BF16 mixed-precision training, and early-stopping callbacks without boilerplate code.

Set up an isolated Python environment and install the core dependencies with CUDA 12.1 support:

python3 -m venv venv
source venv/bin/activate

# Install PyTorch with CUDA 12.1
pip install torch --index-url https://download.pytorch.org/whl/cu121

# Install Sentence-Transformers, Datasets, and Accelerate
pip install -U sentence-transformers datasets accelerate

# Optional: Vietnamese tokenizers for text preprocessing
pip install pyvi underthesea

Verify that PyTorch identifies your GPU before proceeding:

python -c "import torch; print(f'CUDA available: {torch.cuda.is_available()} | Device: {torch.cuda.get_device_name(0)}')"

Configuration

Contrastive learning works best for tailoring dense vector models to niche domains. You feed the model positive (query, passage) pairs or (query, positive, hard negative) triplets. During backpropagation, MultipleNegativesRankingLoss (MNRL) uses every other positive passage in the mini-batch as an implicit negative. A batch size of 32 automatically yields 31 negatives per sample without manual negative curation.

1. Preparing the Training Dataset

We generated 14,000 synthetic questions from our internal text chunks using an LLM, paired them with the source paragraphs, and formatted them into JSON Lines (train_dataset.jsonl):

[{"anchor": "Hồ sơ giảm trừ gia cảnh gồm những giấy tờ gì?", "positive": "Căn cứ theo Thông tư 111/2013/TT-BTC, hồ sơ chứng minh người phụ thuộc đối với con bao gồm bản sao giấy khai sinh..."},
{"anchor": "Quy định thời hạn nộp báo cáo tài chính năm", "positive": "Thời hạn nộp báo cáo tài chính năm chậm nhất là ngày thứ 90 kể từ ngày kết thúc năm dương lịch hoặc năm tài chính..."}]

2. The Training Pipeline

Below is a production-tested training script utilizing SentenceTransformerTrainer, FP16 mixed precision, and in-batch negative contrastive loss:

import os
from datasets import load_dataset
from sentence_transformers import (
    SentenceTransformer,
    SentenceTransformerTrainer,
    SentenceTransformerTrainingArguments,
    losses,
)
from sentence_transformers.training_args import BatchSamplers

# 1. Base model configuration
BASE_MODEL_NAME = "intfloat/multilingual-e5-base"
OUTPUT_DIR = "./vietnamese-domain-embedding-v1"

# 2. Initialize base bi-encoder
model = SentenceTransformer(BASE_MODEL_NAME)

# 3. Load dataset (anchor-positive pairs)
dataset = load_dataset("json", data_files={"train": "train_dataset.jsonl"})
split_dataset = dataset["train"].train_test_split(test_size=0.1, seed=42)
train_ds = split_dataset["train"]
eval_ds = split_dataset["test"]

# 4. Contrastive loss function
train_loss = losses.MultipleNegativesRankingLoss(model)

# 5. Training arguments (optimized for an RTX 4090 / A10G 24GB GPU)
training_args = SentenceTransformerTrainingArguments(
    output_dir=OUTPUT_DIR,
    num_train_epochs=3,
    per_device_train_batch_size=32,  # Larger batches supply richer in-batch negatives
    per_device_eval_batch_size=32,
    learning_rate=2e-5,
    warmup_ratio=0.1,
    fp16=True,                       # Reduces memory footprint and accelerates throughput
    eval_strategy="steps",
    eval_steps=100,
    save_steps=100,
    save_total_limit=2,
    logging_steps=20,
    batch_sampler=BatchSamplers.NO_DUPLICATES,  # Prevents duplicate anchors within the same batch
    load_best_model_at_end=True,
)

# 6. Initialize trainer
trainer = SentenceTransformerTrainer(
    model=model,
    args=training_args,
    train_dataset=train_ds,
    eval_dataset=eval_ds,
    loss=train_loss,
)

# 7. Execute training
trainer.train()

# 8. Export the fine-tuned checkpoint
model.save_pretrained(os.path.join(OUTPUT_DIR, "best_model"))
print("Fine-tuning complete. Checkpoint saved successfully.")

Verification & Monitoring

Do not promote an embedding checkpoint straight to your live vector index based solely on a decreasing training loss. Dense models can fall victim to representation collapse. When this happens, embeddings crowd into a narrow hypersphere cone and lose all discriminative power.

1. Quantitative Evaluation with InformationRetrievalEvaluator

Validate your checkpoint against a held-out test split of 800 unseen queries using ranking metrics: MRR@10, NDCG@10, and Hit@5.

from sentence_transformers.evaluation import InformationRetrievalEvaluator

queries = {
    "q1": "quy định khấu trừ thuế vãng lai 1%",
    "q2": "điều kiện hưởng chế độ thai sản cho lao động nam"
}
corpus = {
    "d1": "Thông tư hướng dẫn về thuế giá trị gia tăng đối với hoạt động xây dựng vãng lai...",
    "d2": "Luật Bảo hiểm xã hội quy định lao động nam có vợ sinh con được nghỉ từ 5 đến 14 ngày làm việc...",
    "d3": "Quy trình bảo dưỡng máy móc định kỳ trong nhà máy thép..."
}
relevant_docs = {
    "q1": {"d1"},
    "q2": {"d2"}
}

evaluator = InformationRetrievalEvaluator(
    queries=queries,
    corpus=corpus,
    relevant_docs=relevant_docs,
    name="vietnamese-ir-benchmark",
    mrr_at_k=[5, 10],
    ndcg_at_k=[10],
    accuracy_at_k=[1, 5]
)

results = evaluator(model)
print(f"MRR@10: {results['vietnamese-ir-benchmark_cosine_mrr@10']:.4f}")
print(f"NDCG@10: {results['vietnamese-ir-benchmark_cosine_ndcg@10']:.4f}")

2. Six Months of Production Observations

After re-indexing our Qdrant vector database with the adapted model, we analyzed 120,000 real production queries:

  • Retrieval Recall: Top-5 hit rate climbed from 61.4% to 89.2% on Vietnamese legal phrasing.
  • Hallucination Rates: Audited outputs flagged for factual inaccuracies dropped by 44%. Feeding accurate excerpts directly reduced generative error.
  • Vector Discrimination: Off-the-shelf multilingual models squeezed cosine scores into a cramped band between 0.78 and 0.88. The fine-tuned bi-encoder restored separation: irrelevant passages hovered around 0.32, while actual answers scored 0.85 and above. This variance gave us an unambiguous cutoff threshold at 0.72 to handle out-of-scope inquiries cleanly.

Keep your embeddings healthy over time. Track cosine drift weekly and log queries that receive user downvotes. Running a scheduled fine-tuning run every quarter incorporates newly released regulations and keeps your retrieval engine reliable.

Share:
Tags: