YuITC/Vietnamese-Legal-Doc-Retrieval

Vietnamese Legal Document Retrieval System

  

A retrieval system specifically designed for Vietnamese legal documents using fine-tuned SBERT (Sentence-BERT) technology.

📌 Overview

This project implements a retrieval system for retrieving relevant Vietnamese legal documents based on user queries. The system uses a fine-tuned multilingual BERT model to encode legal queries and documents into a semantic vector space, allowing for retrieval based on meaning rather than just keyword matching.

[image]

🔑 Key features

• —Step-by-step notebook for understanding.
• —Fine-tuned SBERT model specialized for Vietnamese legal document retrieval.
• —FAISS indexing for efficient vector search.
• —Evaluation based on MTEB.
• —Interactive web interface for quick legal document search.
• —High-performance retrieval of relevant legal passages.

🛠️ Installation & Usage

# Install dependencies
conda install pytorch torchvision torchaudio pytorch-cuda=12.1 -c pytorch -c nvidia
conda install faiss-gpu=1.9.0 -c pytorch -c nvidia
pip install -r requirements.txt

# Running the Application
python main.py

The application will start a local web server with the Gradio interface, allowing you to enter legal queries and retrieve relevant documents.

📂 Project Structure

Vietnamese-Legal-Doc-Retrieval/
├── assets/                   # Visual assets for documentation 
│   └── gradio_demo.png       # Screenshot of the Gradio demo interface
├── cache/                    # Cached model files
│   └── VN-legalDocs-SBERT/   # Cached BERT model files
├── data/                     # Dataset files
│   ├── original/             # Original downloaded dataset
│   │   ├── corpus.csv        # Raw corpus documents
│   │   ├── train_split.csv   # Training data
│   │   ├── val_split.csv     # Validation data
│   │   └── ...
│   ├── processed/            # Processed dataset files
│   │   ├── corpus_data.parquet  # Processed corpus for embedding
│   │   ├── train_data.parquet  # Processed training data
│   │   └── test_data.parquet   # Processed test data
│   └── retrieval/            # Files for retrieval system
│       └── legal_faiss.index # FAISS index for fast vector search
├── models/                   # Trained model files
│   └── VN-legalDocs-SBERT/   # Fine-tuned BERT model for legal documents
│       ├── model.safetensors # Model weights
│       ├── config.json       # Model configuration
│       └── checkpoint-*/     # Training checkpoints
├── results/                  # Evaluation results
├── Dockerfile                # Docker configuration for deployment
├── main.py                   # Main application entry point
├── requirements.txt          # Python dependencies
├── settings.py               # Configuration settings
└── step_*_*.ipynb            # Jupyter notebooks for each step of the process

💾 Dataset

The system is trained on a Vietnamese legal document corpus containing:

• —Legal texts from various domains
• —Query-document pairs for training and evaluation
• —Processed and structured for semantic search training

The dataset is available on Hugging Face (modified by me, the base dataset is cited below).

📊 Model Training Process

The project follows a systematic approach to build the retrieval system:

1. 1.Data Preparation (step_01_Prepare_Data.ipynb):
2. 2.Processes raw legal documents
3. 3.Creates query-document pairs for training
4. 4.Formats data for the embedding model

1. 1.SBERT Fine-tuning (step_02_Finetune_SBERT.ipynb):
2. 2.Fine-tunes a multilingual BERT model with legal document pairs
3. 3.Uses CachedMultipleNegativesRankingLoss for training
4. 4.Optimizes for semantic similarity in legal context

1. 1.Evaluation (step_03_Eval_with_MTEB.ipynb):
2. 2.Evaluates model performance using retrieval metrics
3. 3.Compares with baseline models

1. 1.Retrieval System Setup (step_04_Retrieval.ipynb):
2. 2.Creates FAISS index from document embeddings
3. 3.Implements efficient search functionality
4. 4.Prepares for deployment

🔍 Usage Examples

The system accepts natural language queries in Vietnamese related to legal topics. Example queries:

• —"Tội xúc phạm danh dự?" (Crimes against honor?)
• —"Quyền lợi của người lao động?" (Rights of workers?)
• —"Thủ tục đăng ký kết hôn?" (Marriage registration procedures?)

🧪 Performance

The fine-tuned model was evaluated using the MTEB benchmark on the BKAILegalDocRetrieval dataset. Key results:

• —NDCG@k (Normalized Discounted Cumulative Gain) Measures ranking quality by evaluating the relevance of results with logarithmic position-based discounting.
• —MAP@k (Mean Average Precision) Computes the average precision for each query up to rank k—precision at each relevant retrieved document—then averages across all queries.
• —Recall@k The proportion of all relevant documents that are retrieved in the top k results.
• —Precision@k The proportion of the top k retrieved documents that are relevant.
• —MRR@k (Mean Reciprocal Rank) The average of the reciprocal of the rank position of the first relevant document across all queries.

The model significantly outperforms baseline retrieval methods, with the main evaluation score (NDCG@10) reaching 60.4%, demonstrating strong performance on Vietnamese legal document retrieval tasks.

🐳 Docker Deployment

The project includes a Docker configuration for easy deployment. The Docker image is built on continuumio/miniconda3 and includes GPU support via PyTorch CUDA and FAISS-GPU.

# Build the Docker image
docker build -t vietnamese-legal-retrieval .

# Run the container
docker run -p 7860:7860 vietnamese-legal-retrieval

The container:

• —Uses Python 3.10 with CUDA 12.1 support
• —Installs required dependencies from requirements.txt
• —Exposes port 7860 for the Gradio web interface
• —Sets proper environment variables for security and performance
• —Runs as a non-root user for enhanced security

You can access the web interface by navigating to http://localhost:7860 after starting the container.

📜 License

This project is licensed under the MIT License – feel free to modify and distribute it as needed.

🤝 Acknowledgments

Thanks for:

• —BKAI Legal Retrieval Dataset for the original data
• —Sentence Transformers library for the embedding model architecture
• —Hugging Face for hosting the model and dataset

If you find this project useful, consider ⭐️ starring the repository or contributing to further improvements!

📬 Contact

For any questions or collaboration opportunities, feel free to reach out:

📧 Email: tainguyenphu2502@gmail.com