kietnt0603/nrk-legal-bkai

This is a sentence-transformers model fine-tuned for Vietnamese Legal retrieval, optimized for use in RAG (Retrieval-Augmented Generation) pipelines. It maps Vietnamese legal queries and documents to a 768-dimensional dense vector space, enabling semantic search over Vietnamese legal corpora.

The model was trained with MatryoshkaLoss combined with MultipleNegativesRankingLoss, allowing embeddings to be truncated to smaller dimensions (e.g., 768, 512, 256, 128, 64) with minimal performance degradation.

Model Details

Usage

Direct Usage (Sentence Transformers)

from sentence_transformers import SentenceTransformer

model = SentenceTransformer("kietnt0603/nrk-legal-bkai")

queries = [
    "Mức phạt khi điều khiển xe ô tô không có giấy phép lái xe",
    "Quy định về thời gian nghỉ phép hàng năm của người lao động",
]
documents = [
    "Điều 21. Xử phạt người điều khiển xe ô tô vi phạm quy định...",
    "Điều 113. Nghỉ hằng năm. Người lao động làm việc đủ 12 tháng...",
]

query_embeddings = model.encode(queries)
doc_embeddings = model.encode(documents)
similarities = model.similarity(query_embeddings, doc_embeddings)
print(similarities)

Using Truncated Embeddings (Matryoshka)

from sentence_transformers import SentenceTransformer

model = SentenceTransformer("kietnt0603/nrk-legal-bkai")
model.truncate_dim = 256  # Use 256-dim instead of full 768-dim

query_embedding = model.encode("Quy định về thuế thu nhập cá nhân")
doc_embeddings = model.encode(documents)
# Faster search with minimal quality loss

Evaluation

Evaluated on another-symato/VMTEB-Zalo-legel-retrieval-wseg benchmark (test split: 640 queries, 61,425 legal documents).

Per-Dimension Performance (Matryoshka)

Comparison with Foundation & Community Models

Training

• Base Model: bkai-foundation-models/vietnamese-bi-encoder
• Training Data: Vietnamese legal retrieval triplets from Zalo AI Legal competition (~52K query-document-negative triplets with hard negatives)
• Loss Function: MatryoshkaLoss wrapping MultipleNegativesRankingLoss
• Epochs: 3
• Learning Rate: 2e-05
• Batch Size: 256
• FP16: True
• Best Model Selection: Lowest eval loss (load_best_model_at_end=True)
• Framework: Sentence Transformers v3.0+

Training Details

• Training samples: 49,083
• Eval samples: 2,584
• Total steps: 3,840
• Best eval loss: 2.8625
• Training time: 220 minutes

References

• Kusupati et al., 2024 — Matryoshka Representation Learning
• Henderson et al., 2017 — Efficient Natural Language Response Suggestion for Smart Reply
• Reimers & Gurevych, 2019 — Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks

Citation

@misc{nrk-legal-bkai,
  title={kietnt0603/nrk-legal-bkai: A Vietnamese Legal Text Embedding Model},
  author={Nguyen Tuan Kiet},
  year={2025},
  publisher={Hugging Face},
}