bert-base-multilingual-cased-finetuned-VNLegalDocs

This is a sentence-transformers model finetuned from google-bert/bert-base-multilingual-cased on the BKAI Vietnamese Legal Documents retrieval dataset. It maps sentences & paragraphs to a 768-dimensional dense vector space and can be used for semantic textual similarity, semantic search, paraphrase mining, text classification, clustering, and more. It also achieves an NDCG@10 score of 0.60389.

---

1. Model Details

1.1. Model Description

• Model Type: Sentence Transformer
• Base model: google-bert/bert-base-multilingual-cased
• Maximum Sequence Length: 512 tokens
• Output Dimensionality: 768 dimensions
• Similarity Function: Cosine Similarity
• Training Dataset: BKAI Vietnamese Legal Documents retrieval dataset
• Language: Vietnamese
• License: MIT

1.2. Model Sources

• Documentation: Sentence Transformers Documentation
• Hugging Face: Sentence Transformers on Hugging Face

1.3. Full Model Architecture

SentenceTransformer(
  (0): Transformer({'max_seq_length': 512, 'do_lower_case': False}) with Transformer model: BertModel 
  (1): Pooling({'word_embedding_dimension': 768, 'pooling_mode_cls_token': False, 'pooling_mode_mean_tokens': True, 'pooling_mode_max_tokens': False, 'pooling_mode_mean_sqrt_len_tokens': False, 'pooling_mode_weightedmean_tokens': False, 'pooling_mode_lasttoken': False, 'include_prompt': True})
)

---

2. Usage

2.1. Direct Usage (Sentence Transformers)

1. First install the Sentence Transformers library:

pip install -U sentence-transformers

2. Then you can load this model and run inference.

from sentence_transformers import SentenceTransformer

# Download from the 🤗 Hub & Run inference
model = SentenceTransformer("google-bert/bert-base-multilingual-cased")
sentences = [
    'Tội xúc phạm danh dự?',
    'Quyền lợi của người lao động?',
    'Thủ tục đăng ký kết hôn?',
]
embeddings = model.encode(sentences)
print(embeddings.shape) # [3, 768]

# Get the similarity scores for the embeddings
similarities = model.similarity(embeddings, embeddings)
print(similarities.shape) # [3, 3]

2.2. Downstream Usage - Fine-tune (Sentence Transformers)

You can finetune this model or its base model (google-bert/bert-base-multilingual-cased) on your own dataset.

Click to expand

1. Install dependencies:

pip install sentence-transformers datasets pandas tqdm

2. Prepare your data:

• Organize your data as sentence‐pairs (e.g. query – relevant document)
• Store in Parquet or CSV with columns text_0, text_1

import pandas as pd

df = pd.DataFrame([
    {"text_0": "What is civil procedure?", "text_1": "Civil procedure governs how legal cases are processed."},
    {"text_0": "Define contract law",    "text_1": "Contract law deals with agreements between parties."},
    # …
])
df.to_parquet("data/train.parquet", index=False)

3. Declare important hyperparameters:

MODEL_ID    = "YuITC/bert-base-multilingual-cased-finetuned-VNLegalDocs"
MODEL_NAME  = "bert-base-multilingual-cased"
CACHE_DIR   = "./cache"
OUTPUT_DIR  = "./output"
MAX_SEQ_LEN = 512
EPOCHS      = 5
LR          = 3e-5
BATCH_SIZE  = 128
DEVICE      = "cuda"  # or "cpu"  

4. Build the model & loss:

from sentence_transformers import SentenceTransformer
from sentence_transformers.models import Transformer, Pooling
from sentence_transformers.losses import CachedMultipleNegativesRankingLoss

# Load base transformer + pooling
embedding_model = Transformer(MODEL_ID, max_seq_length=MAX_SEQ_LEN, cache_dir=CACHE_DIR)
pooling_model   = Pooling(embedding_model.get_word_embedding_dimension(),
                          pooling_mode_mean_tokens=True)
model = SentenceTransformer(modules=[embedding_model, pooling_model],
                            device=DEVICE, cache_folder=CACHE_DIR)

# Use multiple negatives ranking loss
loss = CachedMultipleNegativesRankingLoss(model=model)

5. Start fine-tuning:

from sentence_transformers import SentenceTransformerTrainer, SentenceTransformerTrainingArguments
from datasets import Dataset
import pandas as pd

# Load your training DataFrame
df_train = pd.read_parquet("data/train.parquet")
train_ds = Dataset.from_pandas(df_train)

# Training arguments
args = SentenceTransformerTrainingArguments(output_dir=OUTPUT_DIR,
                                            num_train_epochs=EPOCHS,
                                            per_device_train_batch_size=BATCH_SIZE,
                                            learning_rate=LR, warmup_ratio=0.1, fp16=True, logging_steps=100)

trainer = SentenceTransformerTrainer(model=model, args=args, train_dataset=train_ds, loss=loss)

# Start fine-tuning
trainer.train()

6. Load fine-tuned model:

model      = SentenceTransformer(OUTPUT_DIR, device=DEVICE)
embeddings = model.encode(["Your query here"], convert_to_tensor=True)

---

3. Training Details

3.1. Training Dataset: BKAI Vietnamese Legal Documents

• Size: 99,580 training samples (after preprocess)

• Columns: text_0 and text_1

• Approximate statistics based on the first 1000 samples:

  Column	type	min tokens	mean tokens	max tokens
  text_0	string	8	25.64	58
  text_1	string	13	278.08	512

• Loss: CachedMultipleNegativesRankingLoss

3.2. Training Hyperparameters

• per_device_train_batch_size: 128
• learning_rate: 3e-05
• num_train_epochs: 5
• warmup_ratio: 0.1
• fp16: True
• batch_sampler: no_duplicates

3.3. Training Logs

3.4. Training Device Configuraion:

• GPU: 01 x GPU Nvidia RTX A4000 16GB, 6144 CUDA cores
• CPU: 12 Core vCPU AMD EPYC 7K62, 48GB RAM
• Training time: 3:32:33(s)

3.5. Framework Versions

• Python: 3.10.16
• Sentence Transformers: 4.1.0
• Transformers: 4.51.3
• PyTorch: 2.5.1
• Accelerate: 1.6.0
• Datasets: 2.21.0
• Tokenizers: 0.21.1

---

4. Evaluation with MTEB

Click to expand

1. Declare an instance of AbsTaskRetrieval class:

class BKAILegalDocRetrievalTask(AbsTaskRetrieval):
    # Metadata definition used by MTEB benchmark
    metadata = TaskMetadata(name='BKAILegalDocRetrieval',
                            description='',
                            reference='https://github.com/embeddings-benchmark/mteb/blob/main/docs/adding_a_dataset.md',
                            type='Retrieval',
                            category='s2p',
                            modalities=['text'],
                            eval_splits=['test'],
                            eval_langs=['vi'],
                            main_score='ndcg_at_10',
                            other_scores=['recall_at_10', 'precision_at_10', 'map'],
                            dataset={
                                'path'    : 'data',
                                'revision': 'd4c5a8ba10ae71224752c727094ac4c46947fa29',
                            },
                            date=('2012-01-01', '2020-01-01'),
                            form='Written',
                            domains=['Academic', 'Non-fiction'],
                            task_subtypes=['Scientific Reranking'],
                            license='cc-by-nc-4.0',
                            annotations_creators='derived',
                            dialect=[],
                            text_creation='found',
                            bibtex_citation=''
    )

    data_loaded = True # Flag

    def __init__(self, **kwargs):
        super().__init__(**kwargs)

        self.corpus        = {}
        self.queries       = {}
        self.relevant_docs = {}

        shared_corpus = {}
        for _, row in data['corpus'].iterrows():
            shared_corpus[f"c{row['cid']}"] = {
                'text': row['text'],
                '_id' : row['cid']
            }
            
        for split in ['train', 'test']:
            self.corpus[split]        = shared_corpus
            self.queries[split]       = {}
            self.relevant_docs[split] = {}

        for split in ['train', 'test']:
            for _, row in data[split].iterrows():
                qid, cids = row['qid'], row['cid']
                
                qid_str   = f'q{qid}'
                cids_str  = [f'c{cid}' for cid in cids]
                
                self.queries[split][qid_str] = row['question']
                
                if qid_str not in self.relevant_docs[split]:
                    self.relevant_docs[split][qid_str] = {}
                    
                for cid_str in cids_str:
                    self.relevant_docs[split][qid_str][cid_str] = 1
            
        self.data_loaded = True

2. Start evaluate by:

fine_tuned_model = SentenceTransformer(OUTPUT_DIR, device=DEVICE)

custom_task = BKAILegalDocRetrievalTask()
evaluation  = MTEB(tasks=[custom_task])
evaluation.run(fine_tuned_model, batch_size=BATCH_SIZE)

---

5. Citation

Sentence Transformers

@inproceedings{reimers-2019-sentence-bert,
    title = "Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks",
    author = "Reimers, Nils and Gurevych, Iryna",
    booktitle = "Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing",
    month = "11",
    year = "2019",
    publisher = "Association for Computational Linguistics",
    url = "https://arxiv.org/abs/1908.10084",
}

CachedMultipleNegativesRankingLoss

@misc{gao2021scaling,
    title={Scaling Deep Contrastive Learning Batch Size under Memory Limited Setup},
    author={Luyu Gao and Yunyi Zhang and Jiawei Han and Jamie Callan},
    year={2021},
    eprint={2101.06983},
    archivePrefix={arXiv},
    primaryClass={cs.LG}
}