kietnt0603/nrk-legal-bkai

Name: nrk-legal-bkai
Availability: InStock

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

Attribute	Value
Base Model	bkai-foundation-models/vietnamese-bi-encoder
Embedding Dimension	768
Max Sequence Length	256
Similarity Function	Cosine Similarity
Training Loss	MatryoshkaLoss + MultipleNegativesRankingLoss
Matryoshka Dimensions	[768, 512, 256, 128, 64]
Language	Vietnamese

Usage

Direct Usage (Sentence Transformers)

python

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)

python

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)

Dimension	Accuracy@1	Accuracy@10	NDCG@3	NDCG@5	NDCG@10	MRR@3	MRR@5	MRR@10	MAP@100
768	0.6719	0.9406	0.7918	0.8072	0.8161	0.7630	0.7715	0.7751	0.7767
512	0.6719	0.9328	0.7879	0.8055	0.8120	0.7594	0.7693	0.7719	0.7740
256	0.6687	0.9313	0.7848	0.8023	0.8099	0.7568	0.7665	0.7697	0.7716
128	0.6641	0.9328	0.7735	0.7914	0.8049	0.7479	0.7578	0.7633	0.7649
64	0.6062	0.8969	0.7077	0.7382	0.7530	0.6833	0.7004	0.7066	0.7100
hybrid	0.6234	0.9172	0.7166	0.7474	0.7667	0.6937	0.7108	0.7187	0.7233

Comparison with Foundation & Community Models

Model	Type	NDCG@3	NDCG@5	NDCG@10	MRR@3	MRR@5	MRR@10
nrk-legal-large	dense	0.898156	0.907371	0.911489	0.877344	0.882578	0.884320
nrk-legal-base	dense	0.868133	0.878763	0.882884	0.840885	0.846979	0.848726
AITeamVN/Vietnamese_Embedding	dense	0.821426	0.836150	0.848246	0.798958	0.807083	0.812062
nrk-legal-bkai	dense	0.791849	0.807247	0.816148	0.763021	0.771536	0.775099
nrk-legal-base	hybrid	0.762795	0.785115	0.805529	0.741667	0.753932	0.762475
nrk-legal-large	hybrid	0.764636	0.785280	0.805268	0.741406	0.752891	0.761298
nrk-legal-small	dense	0.737869	0.764900	0.781334	0.709896	0.724896	0.731839
nrk-legal-small	hybrid	0.730744	0.750316	0.769867	0.705990	0.717005	0.725277
nrk-legal-bkai	hybrid	0.716627	0.747421	0.766654	0.693750	0.710781	0.718746
bkai-foundation-models/vietnamese-bi-encoder	dense	0.732400	0.750763	0.766062	0.708333	0.718724	0.725125
AITeamVN/Vietnamese_Embedding	hybrid	0.706210	0.726386	0.747192	0.681771	0.693099	0.701738
google/embeddinggemma-300m	dense	0.691310	0.717132	0.733838	0.661198	0.675573	0.682778
bkai-foundation-models/vietnamese-bi-encoder	hybrid	0.677787	0.705491	0.730162	0.651562	0.666953	0.677379
huyydangg/DEk21_hcmute_embedding	hybrid	0.679665	0.699168	0.724153	0.649219	0.660156	0.670600
BAAI/bge-m3	dense	0.672449	0.698202	0.722927	0.646094	0.660391	0.670585
huyydangg/DEk21_hcmute_embedding	dense	0.667296	0.694999	0.713940	0.638021	0.653411	0.661366
BAAI/bge-m3	hybrid	0.650034	0.679015	0.700087	0.621354	0.637448	0.646053
google/embeddinggemma-300m	hybrid	0.641626	0.669535	0.689832	0.609896	0.625521	0.633941
hiieu/halong_embedding	hybrid	0.611511	0.643377	0.667847	0.585677	0.603177	0.613400
intfloat/multilingual-e5-base	hybrid	0.572895	0.602480	0.634951	0.543490	0.559896	0.573371
hiieu/halong_embedding	dense	0.565380	0.594156	0.619427	0.540885	0.556745	0.567188
intfloat/multilingual-e5-base	dense	0.528159	0.555258	0.585950	0.498958	0.514036	0.526622
VoVanPhuc/sup-SimCSE-VietNamese-phobert-base	hybrid	0.515176	0.546712	0.578495	0.486458	0.503958	0.517342
bm25	dense	0.513687	0.548999	0.573213	0.485938	0.505859	0.515820
VoVanPhuc/sup-SimCSE-VietNamese-phobert-base	dense	0.405764	0.431859	0.465527	0.381771	0.396458	0.410246

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

bibtex

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