feat: initial project upload after testing

.gitignore

@@ -0,0 +1,8 @@
+__pycache__/
+.gradio/
+cache/
+data/original/
+models/
+data/
+tmp/
+.env

Dockerfile

@@ -0,0 +1,21 @@
+FROM continuumio/miniconda3
+
+WORKDIR /app
+
+# Get dependencies
+COPY requirements.txt .
+
+RUN conda install -y python=3.10 \
+    pytorch torchvision torchaudio pytorch-cuda=12.1 -c pytorch -c nvidia \
+    faiss-gpu=1.9.0 -c pytorch -c nvidia && \
+    conda clean -afy
+
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Copy the rest of the code
+COPY . /app
+
+# Run the application
+EXPOSE 7860
+
+CMD ["python", "main.py"]

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Nguyen Phu Tai
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -6,6 +6,10 @@ colorTo: pink
 sdk: docker
 pinned: false
 short_description: Fine-tuned Retrieval System for Vietnamese Legal Documents
+models:
+- YuITC/bert-base-multilingual-cased-finetuned-VNLegalDocs
+datasets:
+- YuITC/Vietnamese-Legal-Doc-Retrieval-Data
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

main.py

@@ -0,0 +1,65 @@
+import os
+import numpy as np
+import pandas as pd
+import gradio as gr
+
+import faiss
+from sentence_transformers import SentenceTransformer
+from settings import OUTPUT_DIR, DEVICE
+os.environ['WANDB_DISABLED'] = 'true'
+
+    
+fine_tuned_model = SentenceTransformer(OUTPUT_DIR, device=DEVICE)
+passages         = pd.read_parquet('data/processed/corpus_data.parquet')['text'].tolist()
+legal_index      = faiss.read_index('data/retrieval/legal_faiss.index')
+
+def retrieval(emb_model, query, index, top_k=10):
+    q_emb = emb_model.encode(
+        query, 
+        convert_to_numpy=True, 
+        normalize_embeddings=True,
+    ).astype(np.float32).reshape(1, -1)
+    
+    scores, indices = index.search(q_emb, top_k)  # shape: (1, top_k)
+    
+    cand_idxs   = indices[0]
+    cand_scores = scores[0]
+    cand_texts  = [passages[i] for i in cand_idxs]
+
+    results = [{
+        'index': int(cand_idxs[i]),
+        'score': float(cand_scores[i]),
+        'text': cand_texts[i]
+    } for i in range(len(cand_idxs))]
+    
+    return results
+
+def get_results(query, top_k):
+    hits = retrieval(fine_tuned_model, query, legal_index, top_k=top_k)
+    
+    result = ""
+    for rank, h in enumerate(hits, start=1):
+        result += f"[Kết quả {rank} - Độ tin cậy={h['score']:.4f}]\n\n{h['text']}\n{'-'*100}\n"
+    return result
+
+    
+demo = gr.Interface(
+    'huggingface/YuITC/bert-base-multilingual-cased-finetuned-VNLegalDocs',
+    fn=get_results,
+    inputs=[
+        gr.Textbox(lines=2, placeholder='Nhập câu hỏi pháp lý của bạn...', label='Câu hỏi'),
+        gr.Slider(minimum=5, maximum=20, value=10, step=1, label='Số lượng kết quả'),
+    ],
+    outputs=gr.Textbox(lines=20, label='Kết quả'),
+    title='Vietnamese Legal Document Retrieval System',
+    description='🔍 Nhập câu hỏi pháp lý của bạn bằng tiếng Việt để nhận các đoạn văn bản pháp luật liên quan.',
+    examples=[
+        ['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?'],
+    ],
+    flagging_mode='never'
+)
+
+if __name__ == '__main__':
+    demo.launch(share=True)

requirements.txt

@@ -0,0 +1,12 @@
+# !conda install pytorch torchvision torchaudio pytorch-cuda=12.1 -c pytorch -c nvidia
+# !conda install faiss-gpu=1.9.0 -c pytorch -c nvidia
+
+transformers
+sentence-transformers
+accelerate
+datasets
+mteb
+tqdm
+pandas
+gradio
+huggingface-hub

results/no_model_name_available/no_revision_available/BKAILegalDocRetrieval.json

@@ -0,0 +1,158 @@
+{
+  "dataset_revision": "d4c5a8ba10ae71224752c727094ac4c46947fa29",
+  "task_name": "BKAILegalDocRetrieval",
+  "mteb_version": "1.38.1",
+  "scores": {
+    "test": [
+      {
+        "ndcg_at_1": 0.42425,
+        "ndcg_at_3": 0.53937,
+        "ndcg_at_5": 0.57387,
+        "ndcg_at_10": 0.60389,
+        "ndcg_at_20": 0.6216,
+        "ndcg_at_100": 0.63894,
+        "ndcg_at_1000": 0.64436,
+        "map_at_1": 0.40328,
+        "map_at_3": 0.50314,
+        "map_at_5": 0.52297,
+        "map_at_10": 0.53608,
+        "map_at_20": 0.54136,
+        "map_at_100": 0.54418,
+        "map_at_1000": 0.54442,
+        "recall_at_1": 0.40328,
+        "recall_at_3": 0.62323,
+        "recall_at_5": 0.70466,
+        "recall_at_10": 0.79407,
+        "recall_at_20": 0.86112,
+        "recall_at_100": 0.94805,
+        "recall_at_1000": 0.98787,
+        "precision_at_1": 0.42425,
+        "precision_at_3": 0.22147,
+        "precision_at_5": 0.15119,
+        "precision_at_10": 0.08587,
+        "precision_at_20": 0.04687,
+        "precision_at_100": 0.01045,
+        "precision_at_1000": 0.0011,
+        "mrr_at_1": 0.424183,
+        "mrr_at_3": 0.524672,
+        "mrr_at_5": 0.543372,
+        "mrr_at_10": 0.555102,
+        "mrr_at_20": 0.559556,
+        "mrr_at_100": 0.561719,
+        "mrr_at_1000": 0.561878,
+        "nauc_ndcg_at_1_max": 0.252524,
+        "nauc_ndcg_at_1_std": -0.130263,
+        "nauc_ndcg_at_1_diff1": 0.488176,
+        "nauc_ndcg_at_3_max": 0.298482,
+        "nauc_ndcg_at_3_std": -0.120077,
+        "nauc_ndcg_at_3_diff1": 0.423316,
+        "nauc_ndcg_at_5_max": 0.307625,
+        "nauc_ndcg_at_5_std": -0.110964,
+        "nauc_ndcg_at_5_diff1": 0.419743,
+        "nauc_ndcg_at_10_max": 0.312344,
+        "nauc_ndcg_at_10_std": -0.101157,
+        "nauc_ndcg_at_10_diff1": 0.419576,
+        "nauc_ndcg_at_20_max": 0.31366,
+        "nauc_ndcg_at_20_std": -0.093809,
+        "nauc_ndcg_at_20_diff1": 0.423325,
+        "nauc_ndcg_at_100_max": 0.308888,
+        "nauc_ndcg_at_100_std": -0.091458,
+        "nauc_ndcg_at_100_diff1": 0.428327,
+        "nauc_ndcg_at_1000_max": 0.303777,
+        "nauc_ndcg_at_1000_std": -0.098258,
+        "nauc_ndcg_at_1000_diff1": 0.430885,
+        "nauc_map_at_1_max": 0.238748,
+        "nauc_map_at_1_std": -0.133375,
+        "nauc_map_at_1_diff1": 0.476974,
+        "nauc_map_at_3_max": 0.28179,
+        "nauc_map_at_3_std": -0.124789,
+        "nauc_map_at_3_diff1": 0.435363,
+        "nauc_map_at_5_max": 0.286506,
+        "nauc_map_at_5_std": -0.120112,
+        "nauc_map_at_5_diff1": 0.433864,
+        "nauc_map_at_10_max": 0.288218,
+        "nauc_map_at_10_std": -0.116509,
+        "nauc_map_at_10_diff1": 0.434401,
+        "nauc_map_at_20_max": 0.288517,
+        "nauc_map_at_20_std": -0.114629,
+        "nauc_map_at_20_diff1": 0.435545,
+        "nauc_map_at_100_max": 0.287963,
+        "nauc_map_at_100_std": -0.114142,
+        "nauc_map_at_100_diff1": 0.436281,
+        "nauc_map_at_1000_max": 0.287808,
+        "nauc_map_at_1000_std": -0.114315,
+        "nauc_map_at_1000_diff1": 0.436377,
+        "nauc_recall_at_1_max": 0.238748,
+        "nauc_recall_at_1_std": -0.133375,
+        "nauc_recall_at_1_diff1": 0.476974,
+        "nauc_recall_at_3_max": 0.330773,
+        "nauc_recall_at_3_std": -0.107907,
+        "nauc_recall_at_3_diff1": 0.366506,
+        "nauc_recall_at_5_max": 0.36187,
+        "nauc_recall_at_5_std": -0.080013,
+        "nauc_recall_at_5_diff1": 0.345161,
+        "nauc_recall_at_10_max": 0.399711,
+        "nauc_recall_at_10_std": -0.026693,
+        "nauc_recall_at_10_diff1": 0.318554,
+        "nauc_recall_at_20_max": 0.445634,
+        "nauc_recall_at_20_std": 0.057536,
+        "nauc_recall_at_20_diff1": 0.30652,
+        "nauc_recall_at_100_max": 0.544189,
+        "nauc_recall_at_100_std": 0.325327,
+        "nauc_recall_at_100_diff1": 0.272927,
+        "nauc_recall_at_1000_max": 0.578666,
+        "nauc_recall_at_1000_std": 0.566039,
+        "nauc_recall_at_1000_diff1": 0.23906,
+        "nauc_precision_at_1_max": 0.252524,
+        "nauc_precision_at_1_std": -0.130263,
+        "nauc_precision_at_1_diff1": 0.488176,
+        "nauc_precision_at_3_max": 0.343321,
+        "nauc_precision_at_3_std": -0.090953,
+        "nauc_precision_at_3_diff1": 0.354789,
+        "nauc_precision_at_5_max": 0.356368,
+        "nauc_precision_at_5_std": -0.05169,
+        "nauc_precision_at_5_diff1": 0.308044,
+        "nauc_precision_at_10_max": 0.338907,
+        "nauc_precision_at_10_std": 0.01503,
+        "nauc_precision_at_10_diff1": 0.230763,
+        "nauc_precision_at_20_max": 0.299075,
+        "nauc_precision_at_20_std": 0.08907,
+        "nauc_precision_at_20_diff1": 0.154507,
+        "nauc_precision_at_100_max": 0.148044,
+        "nauc_precision_at_100_std": 0.170043,
+        "nauc_precision_at_100_diff1": 0.008958,
+        "nauc_precision_at_1000_max": 0.011265,
+        "nauc_precision_at_1000_std": 0.110291,
+        "nauc_precision_at_1000_diff1": -0.064328,
+        "nauc_mrr_at_1_max": 0.252492,
+        "nauc_mrr_at_1_std": -0.130181,
+        "nauc_mrr_at_1_diff1": 0.488352,
+        "nauc_mrr_at_3_max": 0.295039,
+        "nauc_mrr_at_3_std": -0.119392,
+        "nauc_mrr_at_3_diff1": 0.445442,
+        "nauc_mrr_at_5_max": 0.298742,
+        "nauc_mrr_at_5_std": -0.114751,
+        "nauc_mrr_at_5_diff1": 0.444417,
+        "nauc_mrr_at_10_max": 0.299598,
+        "nauc_mrr_at_10_std": -0.111479,
+        "nauc_mrr_at_10_diff1": 0.444763,
+        "nauc_mrr_at_20_max": 0.299328,
+        "nauc_mrr_at_20_std": -0.110211,
+        "nauc_mrr_at_20_diff1": 0.44605,
+        "nauc_mrr_at_100_max": 0.298458,
+        "nauc_mrr_at_100_std": -0.110232,
+        "nauc_mrr_at_100_diff1": 0.446632,
+        "nauc_mrr_at_1000_max": 0.298311,
+        "nauc_mrr_at_1000_std": -0.110415,
+        "nauc_mrr_at_1000_diff1": 0.446697,
+        "main_score": 0.60389,
+        "hf_subset": "default",
+        "languages": [
+          "vi"
+        ]
+      }
+    ]
+  },
+  "evaluation_time": 3061.7869832515717,
+  "kg_co2_emissions": null
+}

results/no_model_name_available/no_revision_available/model_meta.json

@@ -0,0 +1 @@
+{"name": "no_model_name_available", "revision": "no_revision_available", "release_date": null, "languages": [], "n_parameters": null, "memory_usage_mb": null, "max_tokens": null, "embed_dim": null, "license": null, "open_weights": true, "public_training_code": null, "public_training_data": null, "framework": ["Sentence Transformers"], "reference": null, "similarity_fn_name": "cosine", "use_instructions": null, "training_datasets": null, "adapted_from": null, "superseded_by": null, "is_cross_encoder": null, "modalities": ["text"], "loader": null}

settings.py

@@ -0,0 +1,44 @@
+import os
+import torch
+import random
+import numpy as np
+
+
+# Data settings
+os.makedirs('data', exist_ok=True)
+os.makedirs('data/original', exist_ok=True)
+os.makedirs('data/processed', exist_ok=True)
+os.makedirs('data/retrieval', exist_ok=True)
+
+
+# Model settings
+MODEL_ID   = 'google-bert/bert-base-multilingual-cased'
+MODEL_NAME = 'VN-legalDocs-SBERT'
+
+CACHE_DIR  = f"cache/{MODEL_NAME}"
+OUTPUT_DIR = f"models/{MODEL_NAME}"
+
+os.makedirs(CACHE_DIR, exist_ok=True)
+os.makedirs(OUTPUT_DIR, exist_ok=True)
+
+
+# Reproducibility
+SEED = 42
+random.seed(SEED)
+np.random.seed(SEED)
+torch.manual_seed(SEED)
+torch.cuda.manual_seed_all(SEED)
+
+# Reproducibility: deterministic=True, benchmark=False
+# Optimize inference/training speed: deterministic=False, benchmark=True
+torch.backends.cudnn.deterministic = False
+torch.backends.cudnn.benchmark     = True
+
+
+# Hyperparameters
+MAX_SEQ_LEN = 512
+EPOCHS      = 5
+LR          = 3e-5
+BATCH_SIZE  = 128
+DEVICE      = 'cuda' if torch.cuda.is_available() else 'cpu'
+print(f"Using device: {DEVICE}")

step_01_Prepare_Data.ipynb

@@ -0,0 +1,411 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "29a91458",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Using device: cuda\n"
+     ]
+    }
+   ],
+   "source": [
+    "!python settings.py"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "97c0ec5c",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "import zipfile\n",
+    "import requests\n",
+    "import pandas as pd"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "f7b1ed51",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Download the dataset\n",
+    "url      = 'https://huggingface.co/datasets/tmnam20/BKAI-Legal-Retrieval/resolve/main/archive.zip'\n",
+    "zip_path = 'data/original/archive.zip'\n",
+    "\n",
+    "response = requests.get(url)\n",
+    "with open(zip_path, 'wb') as f:\n",
+    "    f.write(response.content)\n",
+    "\n",
+    "with zipfile.ZipFile(zip_path, 'r') as zip_ref:\n",
+    "    zip_ref.extractall('data/original')\n",
+    "    \n",
+    "os.remove(zip_path)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "4fe0c4f8",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Train split data: 89592\n",
+      "Test split data : 29864\n"
+     ]
+    }
+   ],
+   "source": [
+    "corpus_data = pd.read_csv('data/original/corpus.csv')\n",
+    "train_split = pd.read_csv('data/original/train_split.csv')\n",
+    "test_split  = pd.read_csv('data/original/val_split.csv')\n",
+    "\n",
+    "print(f\"Train split data: {len(train_split)}\")\n",
+    "print(f\"Test split data : {len(test_split)}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "6e3fbd6e",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/html": [
+       "<div>\n",
+       "<style scoped>\n",
+       "    .dataframe tbody tr th:only-of-type {\n",
+       "        vertical-align: middle;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe tbody tr th {\n",
+       "        vertical-align: top;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe thead th {\n",
+       "        text-align: right;\n",
+       "    }\n",
+       "</style>\n",
+       "<table border=\"1\" class=\"dataframe\">\n",
+       "  <thead>\n",
+       "    <tr style=\"text-align: right;\">\n",
+       "      <th></th>\n",
+       "      <th>text</th>\n",
+       "      <th>cid</th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>Thông tư này hướng dẫn tuần tra, canh gác bảo ...</td>\n",
+       "      <td>0</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1</th>\n",
+       "      <td>1. Hàng năm trước mùa mưa, lũ, Ủy ban nhân dân...</td>\n",
+       "      <td>1</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>2</th>\n",
+       "      <td>Tiêu chuẩn của các thành viên thuộc lực lượng ...</td>\n",
+       "      <td>2</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>3</th>\n",
+       "      <td>Nhiệm vụ của lực lượng tuần tra, canh gác đê\\n...</td>\n",
+       "      <td>3</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>4</th>\n",
+       "      <td>Phù hiệu của lực lượng tuần tra, canh gác đê\\n...</td>\n",
+       "      <td>4</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "</div>"
+      ],
+      "text/plain": [
+       "                                                text  cid\n",
+       "0  Thông tư này hướng dẫn tuần tra, canh gác bảo ...    0\n",
+       "1  1. Hàng năm trước mùa mưa, lũ, Ủy ban nhân dân...    1\n",
+       "2  Tiêu chuẩn của các thành viên thuộc lực lượng ...    2\n",
+       "3  Nhiệm vụ của lực lượng tuần tra, canh gác đê\\n...    3\n",
+       "4  Phù hiệu của lực lượng tuần tra, canh gác đê\\n...    4"
+      ]
+     },
+     "execution_count": 7,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "corpus_data.head()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "3d32d13a",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/html": [
+       "<div>\n",
+       "<style scoped>\n",
+       "    .dataframe tbody tr th:only-of-type {\n",
+       "        vertical-align: middle;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe tbody tr th {\n",
+       "        vertical-align: top;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe thead th {\n",
+       "        text-align: right;\n",
+       "    }\n",
+       "</style>\n",
+       "<table border=\"1\" class=\"dataframe\">\n",
+       "  <thead>\n",
+       "    <tr style=\"text-align: right;\">\n",
+       "      <th></th>\n",
+       "      <th>question</th>\n",
+       "      <th>context</th>\n",
+       "      <th>cid</th>\n",
+       "      <th>qid</th>\n",
+       "      <th>context_list</th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>Liên đoàn Luật sư Việt Nam là tổ chức xã hội –...</td>\n",
+       "      <td>['“Điều 2. Địa vị pháp lý của Liên đoàn Luật s...</td>\n",
+       "      <td>[142820]</td>\n",
+       "      <td>72600</td>\n",
+       "      <td>[“Điều 2. Địa vị pháp lý của Liên đoàn Luật sư...</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1</th>\n",
+       "      <td>Tên hợp tác xã bị rơi vào trường hợp cấm thì c...</td>\n",
+       "      <td>['Tên hợp tác xã, liên hiệp hợp tác xã\\n1. Tên...</td>\n",
+       "      <td>[27817, 72117]</td>\n",
+       "      <td>147562</td>\n",
+       "      <td>[\"Điều 7. Tên hợp tác xã, liên hiệp hợp tác xã...</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>2</th>\n",
+       "      <td>Tài xế lái xe ô tô khách 50 chỗ ngồi bao lâu t...</td>\n",
+       "      <td>['\"1. Sử dụng lái xe bảo đảm sức khỏe theo tiê...</td>\n",
+       "      <td>[33215, 56201]</td>\n",
+       "      <td>142107</td>\n",
+       "      <td>[\"1. Sử dụng lái xe bảo đảm sức khỏe theo tiêu...</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>3</th>\n",
+       "      <td>Các bước chuẩn bị thủ thuật bó bột Cravate sẽ ...</td>\n",
+       "      <td>['BỘT CRAVATE\\n...\\nIV. CHUẨN BỊ\\n1. Người thự...</td>\n",
+       "      <td>[148158]</td>\n",
+       "      <td>77353</td>\n",
+       "      <td>[BỘT CRAVATE\\n...\\nIV. CHUẨN BỊ\\n1. Người thực...</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>4</th>\n",
+       "      <td>Viên chức Hộ sinh hạng 4 có những nhiệm vụ gì ...</td>\n",
+       "      <td>['Hộ sinh hạng IV - Mã số: V.08.06.16\\n1. Nhiệ...</td>\n",
+       "      <td>[188132]</td>\n",
+       "      <td>113090</td>\n",
+       "      <td>[Hộ sinh hạng IV - Mã số: V.08.06.16\\n1. Nhiệm...</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "</div>"
+      ],
+      "text/plain": [
+       "                                            question  \\\n",
+       "0  Liên đoàn Luật sư Việt Nam là tổ chức xã hội –...   \n",
+       "1  Tên hợp tác xã bị rơi vào trường hợp cấm thì c...   \n",
+       "2  Tài xế lái xe ô tô khách 50 chỗ ngồi bao lâu t...   \n",
+       "3  Các bước chuẩn bị thủ thuật bó bột Cravate sẽ ...   \n",
+       "4  Viên chức Hộ sinh hạng 4 có những nhiệm vụ gì ...   \n",
+       "\n",
+       "                                             context             cid     qid  \\\n",
+       "0  ['“Điều 2. Địa vị pháp lý của Liên đoàn Luật s...        [142820]   72600   \n",
+       "1  ['Tên hợp tác xã, liên hiệp hợp tác xã\\n1. Tên...  [27817, 72117]  147562   \n",
+       "2  ['\"1. Sử dụng lái xe bảo đảm sức khỏe theo tiê...  [33215, 56201]  142107   \n",
+       "3  ['BỘT CRAVATE\\n...\\nIV. CHUẨN BỊ\\n1. Người thự...        [148158]   77353   \n",
+       "4  ['Hộ sinh hạng IV - Mã số: V.08.06.16\\n1. Nhiệ...        [188132]  113090   \n",
+       "\n",
+       "                                        context_list  \n",
+       "0  [“Điều 2. Địa vị pháp lý của Liên đoàn Luật sư...  \n",
+       "1  [\"Điều 7. Tên hợp tác xã, liên hiệp hợp tác xã...  \n",
+       "2  [\"1. Sử dụng lái xe bảo đảm sức khỏe theo tiêu...  \n",
+       "3  [BỘT CRAVATE\\n...\\nIV. CHUẨN BỊ\\n1. Người thực...  \n",
+       "4  [Hộ sinh hạng IV - Mã số: V.08.06.16\\n1. Nhiệm...  "
+      ]
+     },
+     "execution_count": 8,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# 'cid' column: '[1 2 3]'\n",
+    "train_split['cid'] = train_split['cid'].apply(lambda x: [int(i) for i in x[1:-1].split()])\n",
+    "test_split['cid']  = test_split['cid'].apply(lambda x: [int(i) for i in x[1:-1].split()])\n",
+    "\n",
+    "\n",
+    "# Mapping from corpus \n",
+    "mapping = dict(zip(corpus_data['cid'], corpus_data['text']))\n",
+    "\n",
+    "def get_context_list(cid_list):\n",
+    "    return [mapping[cid] for cid in cid_list if cid in mapping]\n",
+    "\n",
+    "train_split['context_list'] = train_split['cid'].apply(get_context_list)\n",
+    "test_split['context_list']  = test_split['cid'].apply(get_context_list)\n",
+    "\n",
+    "train_split.head()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "e0450414",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "430 99 331\n",
+      "question <class 'str'>\n",
+      "context <class 'str'>\n",
+      "cid <class 'list'>\n",
+      "qid <class 'numpy.int64'>\n",
+      "context_list <class 'list'>\n"
+     ]
+    }
+   ],
+   "source": [
+    "# Debug\n",
+    "print(\n",
+    "    len(train_split[train_split['context_list'].apply(len) != train_split['cid'].apply(len)]),\n",
+    "    \n",
+    "    len(\n",
+    "        train_split[\n",
+    "            (train_split['context_list'].apply(len) != train_split['cid'].apply(len)) &\n",
+    "            (train_split['context_list'].apply(len) != 0)\n",
+    "        ]\n",
+    "    ),\n",
+    "    \n",
+    "    len(\n",
+    "        train_split[\n",
+    "            (train_split['context_list'].apply(len) != train_split['cid'].apply(len)) &\n",
+    "            (train_split['context_list'].apply(len) == 0)\n",
+    "        ]\n",
+    "    )\n",
+    ")\n",
+    "\n",
+    "for col in train_split.columns:\n",
+    "    print(col, type(train_split[col][0]))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "id": "fd1eb4a2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Drop invalid data\n",
+    "train_data = train_split.loc[\n",
+    "    ~(train_split['context_list'].apply(len) != train_split['cid'].apply(len)), \n",
+    "    ['question', 'context_list', 'qid', 'cid']\n",
+    "]\n",
+    "\n",
+    "test_data = test_split.loc[\n",
+    "    ~(test_split['context_list'].apply(len) != test_split['cid'].apply(len)), \n",
+    "    ['question', 'context_list', 'qid', 'cid']\n",
+    "]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "id": "3661c9cb",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Train data saved: 89162\n",
+      "Test data saved : 29723\n"
+     ]
+    }
+   ],
+   "source": [
+    "# Save the processed data to parquet files\n",
+    "corpus_data.to_parquet('data/processed/corpus_data.parquet', index=False)\n",
+    "train_data.to_parquet('data/processed/train_data.parquet', index=False)\n",
+    "test_data.to_parquet('data/processed/test_data.parquet', index=False)\n",
+    "\n",
+    "print(f\"Train data saved: {len(train_data)}\")\n",
+    "print(f\"Test data saved : {len(test_data)}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "id": "6382a715",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# # Get demo data\n",
+    "# os.makedirs('data/demo', exist_ok=True)\n",
+    "\n",
+    "# demo_corpus_data = corpus_data.sample(10, random_state=42).reset_index(drop=True)\n",
+    "# demo_train_data  = train_data.sample(10, random_state=42).reset_index(drop=True)\n",
+    "# demo_test_data   = test_data.sample(10, random_state=42).reset_index(drop=True)\n",
+    "\n",
+    "# demo_corpus_data.to_csv('data/demo/demo_corpus_data.csv', index=False)\n",
+    "# demo_train_data.to_csv('data/demo/demo_train_data.csv', index=False)\n",
+    "# demo_test_data.to_csv('data/demo/demo_test_data.csv', index=False)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "legal_doc_retrieval",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.16"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

step_02_Finetune_SBERT.ipynb

@@ -0,0 +1,580 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "24106202",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Using device: cuda\n"
+     ]
+    }
+   ],
+   "source": [
+    "!python settings.py"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "0086aabe",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Using device: cuda\n"
+     ]
+    }
+   ],
+   "source": [
+    "import os\n",
+    "import pandas as pd\n",
+    "from datasets import Dataset\n",
+    "from tqdm.autonotebook import tqdm\n",
+    "\n",
+    "from sentence_transformers import (\n",
+    "    SentenceTransformer,\n",
+    "    SentenceTransformerTrainer,\n",
+    "    SentenceTransformerTrainingArguments,\n",
+    "    SentenceTransformerModelCardData,\n",
+    ")\n",
+    "from sentence_transformers.readers       import InputExample\n",
+    "from sentence_transformers.models        import Transformer, Pooling\n",
+    "from sentence_transformers.losses        import CachedMultipleNegativesRankingLoss\n",
+    "from sentence_transformers.training_args import BatchSamplers\n",
+    "\n",
+    "from settings import MODEL_ID, MODEL_NAME, CACHE_DIR, OUTPUT_DIR, MAX_SEQ_LEN, EPOCHS, LR, BATCH_SIZE, DEVICE\n",
+    "\n",
+    "os.environ['WANDB_DISABLED'] = 'true'"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "3a5cc53d",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "data = {\n",
+    "    'corpus': pd.read_parquet('data/processed/corpus_data.parquet'),\n",
+    "    'train' : pd.read_parquet('data/processed/train_data.parquet'),\n",
+    "    'test'  : pd.read_parquet('data/processed/test_data.parquet')\n",
+    "}\n",
+    "for split in ['train', 'test']:\n",
+    "    data[split]['cid']          = data[split]['cid'].apply(lambda x: x.tolist())\n",
+    "    data[split]['context_list'] = data[split]['context_list'].apply(lambda x: x.tolist())\n",
+    "    \n",
+    "examples = {'train': [], 'test': []}"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "30ebbd40",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/html": [
+       "<div>\n",
+       "<style scoped>\n",
+       "    .dataframe tbody tr th:only-of-type {\n",
+       "        vertical-align: middle;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe tbody tr th {\n",
+       "        vertical-align: top;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe thead th {\n",
+       "        text-align: right;\n",
+       "    }\n",
+       "</style>\n",
+       "<table border=\"1\" class=\"dataframe\">\n",
+       "  <thead>\n",
+       "    <tr style=\"text-align: right;\">\n",
+       "      <th></th>\n",
+       "      <th>question</th>\n",
+       "      <th>context_list</th>\n",
+       "      <th>qid</th>\n",
+       "      <th>cid</th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>Liên đoàn Luật sư Việt Nam là tổ chức xã hội –...</td>\n",
+       "      <td>[“Điều 2. Địa vị pháp lý của Liên đoàn Luật sư...</td>\n",
+       "      <td>72600</td>\n",
+       "      <td>[142820]</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1</th>\n",
+       "      <td>Tên hợp tác xã bị rơi vào trường hợp cấm thì c...</td>\n",
+       "      <td>[\"Điều 7. Tên hợp tác xã, liên hiệp hợp tác xã...</td>\n",
+       "      <td>147562</td>\n",
+       "      <td>[27817, 72117]</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>2</th>\n",
+       "      <td>Tài xế lái xe ô tô khách 50 chỗ ngồi bao lâu t...</td>\n",
+       "      <td>[\"1. Sử dụng lái xe bảo đảm sức khỏe theo tiêu...</td>\n",
+       "      <td>142107</td>\n",
+       "      <td>[33215, 56201]</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>3</th>\n",
+       "      <td>Các bước chuẩn bị thủ thuật bó bột Cravate sẽ ...</td>\n",
+       "      <td>[BỘT CRAVATE\\n...\\nIV. CHUẨN BỊ\\n1. Người thực...</td>\n",
+       "      <td>77353</td>\n",
+       "      <td>[148158]</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>4</th>\n",
+       "      <td>Viên chức Hộ sinh hạng 4 có những nhiệm vụ gì ...</td>\n",
+       "      <td>[Hộ sinh hạng IV - Mã số: V.08.06.16\\n1. Nhiệm...</td>\n",
+       "      <td>113090</td>\n",
+       "      <td>[188132]</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "</div>"
+      ],
+      "text/plain": [
+       "                                            question  \\\n",
+       "0  Liên đoàn Luật sư Việt Nam là tổ chức xã hội –...   \n",
+       "1  Tên hợp tác xã bị rơi vào trường hợp cấm thì c...   \n",
+       "2  Tài xế lái xe ô tô khách 50 chỗ ngồi bao lâu t...   \n",
+       "3  Các bước chuẩn bị thủ thuật bó bột Cravate sẽ ...   \n",
+       "4  Viên chức Hộ sinh hạng 4 có những nhiệm vụ gì ...   \n",
+       "\n",
+       "                                        context_list     qid             cid  \n",
+       "0  [“Điều 2. Địa vị pháp lý của Liên đoàn Luật sư...   72600        [142820]  \n",
+       "1  [\"Điều 7. Tên hợp tác xã, liên hiệp hợp tác xã...  147562  [27817, 72117]  \n",
+       "2  [\"1. Sử dụng lái xe bảo đảm sức khỏe theo tiêu...  142107  [33215, 56201]  \n",
+       "3  [BỘT CRAVATE\\n...\\nIV. CHUẨN BỊ\\n1. Người thực...   77353        [148158]  \n",
+       "4  [Hộ sinh hạng IV - Mã số: V.08.06.16\\n1. Nhiệm...  113090        [188132]  "
+      ]
+     },
+     "execution_count": 4,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "data['train'].head()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "943bf8ce",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "question <class 'str'>\n",
+      "context_list <class 'list'>\n",
+      "qid <class 'numpy.int64'>\n",
+      "cid <class 'list'>\n",
+      "True\n"
+     ]
+    }
+   ],
+   "source": [
+    "# Debug\n",
+    "for col in data['test'].columns:\n",
+    "    print(col, type(data['test'][col][0]))\n",
+    "    \n",
+    "print((data['test']['cid'].apply(len) == data['test']['context_list'].apply(len)).all())"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "2c751cf4",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "509893cf5cfd4a8d9e18bba47561a41c",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Processing train:   0%|          | 0/89162 [00:00<?, ?rows/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "12f4fcee4e4244128d8fb472881862ae",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Processing test:   0%|          | 0/29723 [00:00<?, ?rows/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Training examples: 99580\n"
+     ]
+    }
+   ],
+   "source": [
+    "for split in ['train', 'test']:\n",
+    "    rows = list(data[split].itertuples(index=False))\n",
+    "    \n",
+    "    for row in tqdm(rows, desc=f\"Processing {split}\", unit='rows'):\n",
+    "        q = row.question\n",
+    "        for c in row.context_list:\n",
+    "            examples[split].append(InputExample(texts=[q, c]))\n",
+    "\n",
+    "print(f\"Training examples: {len(examples['train'])}\") # Compare with sum(data['train']['cid'].apply(len))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "aadda6e7",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "embedding_model = Transformer(MODEL_ID, max_seq_length=MAX_SEQ_LEN, cache_dir=CACHE_DIR)\n",
+    "pooling_model   = Pooling(\n",
+    "    embedding_model.get_word_embedding_dimension(), \n",
+    "    pooling_mode_mean_tokens=True\n",
+    ")\n",
+    "\n",
+    "model = SentenceTransformer(\n",
+    "    modules=[embedding_model, pooling_model], device=DEVICE, \n",
+    "    cache_folder=CACHE_DIR,\n",
+    "    model_card_data=SentenceTransformerModelCardData(\n",
+    "        model_id=MODEL_ID, \n",
+    "        model_name=MODEL_NAME, \n",
+    "        language='vi',\n",
+    "        license='mit',\n",
+    "    )\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8967eb55",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Using the `WANDB_DISABLED` environment variable is deprecated and will be removed in v5. Use the --report_to flag to control the integrations used for logging result (for instance --report_to none).\n"
+     ]
+    }
+   ],
+   "source": [
+    "loss = CachedMultipleNegativesRankingLoss(model=model)\n",
+    "\n",
+    "args = SentenceTransformerTrainingArguments(\n",
+    "    output_dir=OUTPUT_DIR,\n",
+    "    num_train_epochs=EPOCHS,\n",
+    "    per_device_train_batch_size=BATCH_SIZE,\n",
+    "    learning_rate=LR,\n",
+    "    warmup_ratio=0.1,\n",
+    "    fp16=True,\n",
+    "    batch_sampler=BatchSamplers.NO_DUPLICATES,\n",
+    "    logging_steps=100\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "8bb935fe",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Using the `WANDB_DISABLED` environment variable is deprecated and will be removed in v5. Use the --report_to flag to control the integrations used for logging result (for instance --report_to none).\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "3d68dc4ff84244488d9de723e68b37ca",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Computing widget examples:   0%|          | 0/1 [00:00<?, ?example/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       "\n",
+       "    <div>\n",
+       "      \n",
+       "      <progress value='3890' max='3890' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
+       "      [3890/3890 3:32:33, Epoch 5/5]\n",
+       "    </div>\n",
+       "    <table border=\"1\" class=\"dataframe\">\n",
+       "  <thead>\n",
+       " <tr style=\"text-align: left;\">\n",
+       "      <th>Step</th>\n",
+       "      <th>Training Loss</th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <td>100</td>\n",
+       "      <td>1.882700</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>200</td>\n",
+       "      <td>0.442800</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>300</td>\n",
+       "      <td>0.356400</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>400</td>\n",
+       "      <td>0.285600</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>500</td>\n",
+       "      <td>0.244500</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>600</td>\n",
+       "      <td>0.224100</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>700</td>\n",
+       "      <td>0.193800</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>800</td>\n",
+       "      <td>0.189400</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>900</td>\n",
+       "      <td>0.143200</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>1000</td>\n",
+       "      <td>0.143200</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>1100</td>\n",
+       "      <td>0.134100</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>1200</td>\n",
+       "      <td>0.131100</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>1300</td>\n",
+       "      <td>0.124900</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>1400</td>\n",
+       "      <td>0.122700</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>1500</td>\n",
+       "      <td>0.124100</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>1600</td>\n",
+       "      <td>0.102800</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>1700</td>\n",
+       "      <td>0.085200</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>1800</td>\n",
+       "      <td>0.085000</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>1900</td>\n",
+       "      <td>0.082000</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>2000</td>\n",
+       "      <td>0.080000</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>2100</td>\n",
+       "      <td>0.082400</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>2200</td>\n",
+       "      <td>0.080200</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>2300</td>\n",
+       "      <td>0.082200</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>2400</td>\n",
+       "      <td>0.063300</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>2500</td>\n",
+       "      <td>0.061500</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>2600</td>\n",
+       "      <td>0.061200</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>2700</td>\n",
+       "      <td>0.058000</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>2800</td>\n",
+       "      <td>0.056600</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>2900</td>\n",
+       "      <td>0.052100</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>3000</td>\n",
+       "      <td>0.054800</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>3100</td>\n",
+       "      <td>0.054700</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>3200</td>\n",
+       "      <td>0.047900</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>3300</td>\n",
+       "      <td>0.044900</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>3400</td>\n",
+       "      <td>0.044000</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>3500</td>\n",
+       "      <td>0.043900</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>3600</td>\n",
+       "      <td>0.044400</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>3700</td>\n",
+       "      <td>0.045700</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>3800</td>\n",
+       "      <td>0.046100</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table><p>"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/plain": [
+       "TrainOutput(global_step=3890, training_loss=0.1604946916084976, metrics={'train_runtime': 12756.5123, 'train_samples_per_second': 39.031, 'train_steps_per_second': 0.305, 'total_flos': 0.0, 'train_loss': 0.1604946916084976, 'epoch': 5.0})"
+      ]
+     },
+     "execution_count": 9,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "def to_frame(ex_list):\n",
+    "    rows = [(ex.texts[0], ex.texts[1]) for ex in ex_list]\n",
+    "    return pd.DataFrame(rows, columns=['text_0', 'text_1'])\n",
+    "\n",
+    "train_ds = Dataset.from_pandas(to_frame(examples['train']))\n",
+    "\n",
+    "trainer = SentenceTransformerTrainer(\n",
+    "    model=model,\n",
+    "    args=args,\n",
+    "    train_dataset=train_ds,\n",
+    "    loss=loss,\n",
+    ")\n",
+    "trainer.train()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f47a01a1",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model.save_pretrained(OUTPUT_DIR)\n",
+    "# model.push_to_hub(\n",
+    "#     repo_id='YuITC/bert-base-multilingual-cased-finetuned-VNLegalDocs', \n",
+    "#     commit_message='Update README.md',\n",
+    "#     exist_ok=True,\n",
+    "#     replace_model_card=False,\n",
+    "#     train_datasets=['tmnam20/BKAI-Legal-Retrieval']\n",
+    "# )"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "legal_doc_retrieval",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.16"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

step_03_Eval_with_MTEB.ipynb

@@ -0,0 +1,479 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "b41fd227",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Using device: cuda\n"
+     ]
+    }
+   ],
+   "source": [
+    "!python settings.py"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "b5fd917b",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "📦 PyTorch version: 2.5.1\n",
+      "🚀 CUDA available : True\n",
+      "🧠 GPU Name       : NVIDIA RTX A4000\n",
+      "📦 FAISS version : 1.9.0\n",
+      "🚀 FAISS is using GPU ✅\n"
+     ]
+    }
+   ],
+   "source": [
+    "import torch\n",
+    "\n",
+    "print(\"📦 PyTorch version:\", torch.__version__)\n",
+    "print(\"🚀 CUDA available :\", torch.cuda.is_available())\n",
+    "if torch.cuda.is_available():\n",
+    "    print(\"🧠 GPU Name       :\", torch.cuda.get_device_name(0))\n",
+    "    \n",
+    "import faiss\n",
+    "\n",
+    "print(\"📦 FAISS version :\", faiss.__version__)\n",
+    "\n",
+    "# Kiểm tra module FAISS-GPU có hoạt động không\n",
+    "try:\n",
+    "    res = faiss.StandardGpuResources()  # Nếu không lỗi là có GPU\n",
+    "    print(\"🚀 FAISS is using GPU ✅\")\n",
+    "except Exception as e:\n",
+    "    print(\"❌ FAISS is NOT using GPU:\", str(e))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "030016c2",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "C:\\Users\\Administrator\\AppData\\Local\\Temp\\2\\ipykernel_648\\3951191562.py:5: TqdmExperimentalWarning: Using `tqdm.autonotebook.tqdm` in notebook mode. Use `tqdm.tqdm` instead to force console mode (e.g. in jupyter console)\n",
+      "  from tqdm.autonotebook import tqdm\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Using device: cuda\n"
+     ]
+    }
+   ],
+   "source": [
+    "import os\n",
+    "import json\n",
+    "import pandas as pd\n",
+    "from pprint import pprint\n",
+    "from tqdm.autonotebook import tqdm\n",
+    "\n",
+    "from sentence_transformers import SentenceTransformer\n",
+    "from mteb import MTEB\n",
+    "from mteb.abstasks.TaskMetadata import TaskMetadata\n",
+    "from mteb.abstasks.AbsTaskRetrieval import AbsTaskRetrieval\n",
+    "\n",
+    "from settings import MODEL_NAME, OUTPUT_DIR, DEVICE, BATCH_SIZE\n",
+    "\n",
+    "os.environ['WANDB_DISABLED'] = 'true'"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "dd3f53a3",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "data = {\n",
+    "    'corpus': pd.read_parquet('data/processed/corpus_data.parquet'),\n",
+    "    'train' : pd.read_parquet('data/processed/train_data.parquet'),\n",
+    "    'test'  : pd.read_parquet('data/processed/test_data.parquet')\n",
+    "}\n",
+    "for split in ['train', 'test']:\n",
+    "    data[split]['cid']          = data[split]['cid'].apply(lambda x: x.tolist())\n",
+    "    data[split]['context_list'] = data[split]['context_list'].apply(lambda x: x.tolist())"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "41ffd5ce",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class BKAILegalDocRetrievalTask(AbsTaskRetrieval):\n",
+    "    # Metadata definition used by MTEB benchmark\n",
+    "    metadata = TaskMetadata(name='BKAILegalDocRetrieval',\n",
+    "                            description='',\n",
+    "                            reference='https://github.com/embeddings-benchmark/mteb/blob/main/docs/adding_a_dataset.md',\n",
+    "                            type='Retrieval',\n",
+    "                            category='s2p',\n",
+    "                            modalities=['text'],\n",
+    "                            eval_splits=['test'],\n",
+    "                            eval_langs=['vi'],\n",
+    "                            main_score='ndcg_at_10',\n",
+    "                            other_scores=['recall_at_10', 'precision_at_10', 'map'],\n",
+    "                            dataset={\n",
+    "                                'path'    : 'data',\n",
+    "                                'revision': 'd4c5a8ba10ae71224752c727094ac4c46947fa29',\n",
+    "                            },\n",
+    "                            date=('2012-01-01', '2020-01-01'),\n",
+    "                            form='Written',\n",
+    "                            domains=['Academic', 'Non-fiction'],\n",
+    "                            task_subtypes=['Scientific Reranking'],\n",
+    "                            license='cc-by-nc-4.0',\n",
+    "                            annotations_creators='derived',\n",
+    "                            dialect=[],\n",
+    "                            text_creation='found',\n",
+    "                            bibtex_citation=''\n",
+    "    )\n",
+    "\n",
+    "    data_loaded = True # Flag\n",
+    "\n",
+    "    def __init__(self, **kwargs):\n",
+    "        super().__init__(**kwargs)\n",
+    "\n",
+    "        self.corpus        = {}\n",
+    "        self.queries       = {}\n",
+    "        self.relevant_docs = {}\n",
+    "\n",
+    "        shared_corpus = {}\n",
+    "        for _, row in data['corpus'].iterrows():\n",
+    "            shared_corpus[f\"c{row['cid']}\"] = {\n",
+    "                'text': row['text'],\n",
+    "                '_id' : row['cid']\n",
+    "            }\n",
+    "            \n",
+    "        for split in ['train', 'test']:\n",
+    "            self.corpus[split]        = shared_corpus\n",
+    "            self.queries[split]       = {}\n",
+    "            self.relevant_docs[split] = {}\n",
+    "\n",
+    "        for split in ['train', 'test']:\n",
+    "            for _, row in data[split].iterrows():\n",
+    "                qid, cids = row['qid'], row['cid']\n",
+    "                \n",
+    "                qid_str   = f'q{qid}'\n",
+    "                cids_str  = [f'c{cid}' for cid in cids]\n",
+    "                \n",
+    "                self.queries[split][qid_str] = row['question']\n",
+    "                \n",
+    "                if qid_str not in self.relevant_docs[split]:\n",
+    "                    self.relevant_docs[split][qid_str] = {}\n",
+    "                    \n",
+    "                for cid_str in cids_str:\n",
+    "                    self.relevant_docs[split][qid_str][cid_str] = 1\n",
+    "            \n",
+    "        self.data_loaded = True"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "8c212fe9",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "fine_tuned_model = SentenceTransformer(OUTPUT_DIR, device=DEVICE)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "aae09322",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "The `batch_size` argument is deprecated and will be removed in the next release. Please use `encode_kwargs = {'batch_size': ...}` to set the batch size instead.\n"
+     ]
+    },
+    {
+     "data": {
+      "text/html": [
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #262626; text-decoration-color: #262626\">───────────────────────────────────────────────── </span><span style=\"font-weight: bold\">Selected tasks </span><span style=\"color: #262626; text-decoration-color: #262626\"> ─────────────────────────────────────────────────</span>\n",
+       "</pre>\n"
+      ],
+      "text/plain": [
+       "\u001b[38;5;235m───────────────────────────────────────────────── \u001b[0m\u001b[1mSelected tasks \u001b[0m\u001b[38;5;235m ─────────────────────────────────────────────────\u001b[0m\n"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\">Retrieval</span>\n",
+       "</pre>\n"
+      ],
+      "text/plain": [
+       "\u001b[1mRetrieval\u001b[0m\n"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">    - BKAILegalDocRetrieval, <span style=\"color: #626262; text-decoration-color: #626262; font-style: italic\">s2p</span>\n",
+       "</pre>\n"
+      ],
+      "text/plain": [
+       "    - BKAILegalDocRetrieval, \u001b[3;38;5;241ms2p\u001b[0m\n"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">\n",
+       "\n",
+       "</pre>\n"
+      ],
+      "text/plain": [
+       "\n",
+       "\n"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "53778754caf4456f8e140cfa58b60709",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Batches:   0%|          | 0/233 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "f9b27ae885fc46ad83f332f222a76381",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Batches:   0%|          | 0/391 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "0b6e38b0d54a4b429db05158604d24a5",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Batches:   0%|          | 0/391 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "20ec5df7261c43a7921abc968cc5e3a6",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Batches:   0%|          | 0/391 [00:02<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "5f365f06d3de4becb965adb801aeee60",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Batches:   0%|          | 0/391 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "a43b764ac83e43aeb754c1e60771fd5c",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Batches:   0%|          | 0/391 [00:02<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "ae46c8f76bc64eac8ca475d13f312875",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Batches:   0%|          | 0/91 [00:02<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/plain": [
+       "[TaskResult(task_name=BKAILegalDocRetrieval, scores=...)]"
+      ]
+     },
+     "execution_count": 7,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "custom_task = BKAILegalDocRetrievalTask()\n",
+    "evaluation  = MTEB(tasks=[custom_task])\n",
+    "evaluation.run(fine_tuned_model, batch_size=BATCH_SIZE)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "004e6930",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Main Evaluation Metrics (Top-K = 10):\n",
+      "{'evaluation_time (s)': 3061.7869832515717,\n",
+      " 'main_score': 0.60389,\n",
+      " 'mrr@10': 0.555102,\n",
+      " 'precision@10': 0.08587,\n",
+      " 'recall@10': 0.79407}\n"
+     ]
+    }
+   ],
+   "source": [
+    "file_path = f\"results/no_model_name_available/no_revision_available/BKAILegalDocRetrieval.json\"\n",
+    "\n",
+    "with open(file_path, 'r', encoding='utf-8') as f:\n",
+    "    eval_data = json.load(f)\n",
+    "\n",
+    "scores = eval_data[\"scores\"][\"test\"][0]\n",
+    "main_metrics = {\n",
+    "    'main_score'         : scores.get('ndcg_at_10'),\n",
+    "    'recall@10'          : scores.get('recall_at_10'),\n",
+    "    'precision@10'       : scores.get('precision_at_10'),\n",
+    "    'mrr@10'             : scores.get('mrr_at_10'),\n",
+    "    'evaluation_time (s)': eval_data.get('evaluation_time')\n",
+    "}\n",
+    "\n",
+    "print('Main Evaluation Metrics (Top-K = 10):')\n",
+    "pprint(main_metrics)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "672ebc32",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "Evaluation Scores by K:\n",
+      "metric     map     mrr    ndcg  precision  recall\n",
+      "k                                                \n",
+      "1       0.4033  0.4242  0.4242     0.4242  0.4033\n",
+      "3       0.5031  0.5247  0.5394     0.2215  0.6232\n",
+      "5       0.5230  0.5434  0.5739     0.1512  0.7047\n",
+      "10      0.5361  0.5551  0.6039     0.0859  0.7941\n",
+      "20      0.5414  0.5596  0.6216     0.0469  0.8611\n",
+      "100     0.5442  0.5617  0.6389     0.0104  0.9480\n",
+      "1000    0.5444  0.5619  0.6444     0.0011  0.9879\n"
+     ]
+    }
+   ],
+   "source": [
+    "metrics = {k: v for k, v in scores.items() if '_at_' in k and not k.startswith('nauc')}\n",
+    "\n",
+    "parsed_metrics = []\n",
+    "for key, value in metrics.items():\n",
+    "    metric, at_k = key.split('_at_')\n",
+    "    parsed_metrics.append({'metric': metric, 'k': int(at_k), 'score': value})\n",
+    "\n",
+    "df_metrics = pd.DataFrame(parsed_metrics).pivot(index='k', columns='metric', values='score')\n",
+    "df_metrics = df_metrics.sort_index()\n",
+    "\n",
+    "print(\"\\nEvaluation Scores by K:\")\n",
+    "print(df_metrics.round(4))"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "legal_doc_retrieval",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.16"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

step_04_Retrieval.ipynb

@@ -0,0 +1,383 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "id": "1195e917",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Using device: cuda\n"
+     ]
+    }
+   ],
+   "source": [
+    "!python settings.py"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "id": "01589fc8",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "import numpy as np\n",
+    "import pandas as pd\n",
+    "from tqdm.autonotebook import tqdm\n",
+    "\n",
+    "import faiss\n",
+    "from sentence_transformers import SentenceTransformer, CrossEncoder\n",
+    "\n",
+    "from settings import OUTPUT_DIR, DEVICE\n",
+    "\n",
+    "os.environ['WANDB_DISABLED'] = 'true'\n",
+    "\n",
+    "from transformers import logging\n",
+    "logging.set_verbosity_error()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "id": "057e852f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# data = {\n",
+    "#     'corpus': pd.read_parquet('data/processed/corpus_data.parquet'),\n",
+    "#     'train' : pd.read_parquet('data/processed/train_data.parquet'),\n",
+    "#     'test'  : pd.read_parquet('data/processed/test_data.parquet')\n",
+    "# }\n",
+    "# for split in ['train', 'test']:\n",
+    "#     data[split]['cid']          = data[split]['cid'].apply(lambda x: x.tolist())\n",
+    "#     data[split]['context_list'] = data[split]['context_list'].apply(lambda x: x.tolist())"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 14,
+   "id": "5634b72a",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "SentenceTransformer(\n",
+       "  (0): Transformer({'max_seq_length': 512, 'do_lower_case': False}) with Transformer model: BertModel \n",
+       "  (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})\n",
+       ")"
+      ]
+     },
+     "execution_count": 14,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "fine_tuned_model = SentenceTransformer(OUTPUT_DIR, device=DEVICE)\n",
+    "fine_tuned_model.half()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 15,
+   "id": "62cc0ead",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "passages          = pd.read_parquet('data/processed/corpus_data.parquet')['text'].tolist()\n",
+    "# corpus_embeddings = fine_tuned_model.encode(\n",
+    "#     passages, \n",
+    "#     batch_size=128,\n",
+    "#     convert_to_numpy=True, \n",
+    "#     normalize_embeddings=True,\n",
+    "#     show_progress_bar=True, \n",
+    "#     device=DEVICE,\n",
+    "# ).astype(np.float32)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 16,
+   "id": "465e8d2a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# d         = corpus_embeddings.shape[1]  # 768\n",
+    "# cpu_index = faiss.IndexFlatIP(d)\n",
+    "\n",
+    "# res       = faiss.StandardGpuResources()\n",
+    "# gpu_index = faiss.index_cpu_to_gpu(res, 0, cpu_index)\n",
+    "# gpu_index.add(corpus_embeddings)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 17,
+   "id": "af365371",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# final_cpu_index = faiss.index_gpu_to_cpu(gpu_index)\n",
+    "# faiss.write_index(final_cpu_index, 'data/retrieval/legal_faiss.index')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 18,
+   "id": "9251d0db",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "legal_index = faiss.read_index('data/retrieval/legal_faiss.index')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 19,
+   "id": "9f54c596",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def retrieval(emb_model, query, index, top_k=10):\n",
+    "    q_emb = emb_model.encode(\n",
+    "        query, \n",
+    "        convert_to_numpy=True, \n",
+    "        normalize_embeddings=True,\n",
+    "    ).astype(np.float32).reshape(1, -1)\n",
+    "    \n",
+    "    scores, indices = index.search(q_emb, top_k)  # shape: (1, top_k)\n",
+    "    \n",
+    "    cand_idxs   = indices[0]\n",
+    "    cand_scores = scores[0]\n",
+    "    cand_texts  = [passages[i] for i in cand_idxs]\n",
+    "\n",
+    "    results = [{\n",
+    "        'index': int(cand_idxs[i]),\n",
+    "        'score': float(cand_scores[i]),\n",
+    "        'text': cand_texts[i]\n",
+    "    } for i in range(len(cand_idxs))]\n",
+    "    \n",
+    "    return results"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 22,
+   "id": "ece21ef6",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[Rank 1] index=76423, score=0.6417\n",
+      "Tội làm nhục người khác\n",
+      "1. Người nào xúc phạm nghiêm trọng nhân phẩm, danh dự của người khác, thì bị phạt cảnh cáo, phạt tiền từ 10.000.000 đồng đến 30.000.000 đồng hoặc phạt cải tạo không giam giữ đến 03 năm.\n",
+      "...\n",
+      "--------------------------------------------------------------------------------\n",
+      "[Rank 2] index=99131, score=0.6155\n",
+      "“Người nào có hành vi xâm phạm danh dự, nhân phẩm của người khác mà gây thiệt hại thì phải bồi thường.”\n",
+      "--------------------------------------------------------------------------------\n",
+      "[Rank 3] index=228550, score=0.5932\n",
+      "i) Điều 353, các khoản 2, 3 và 4 (tội tham ô tài sản); Điều 354, các khoản 2, 3 và 4 (tội nhận hối lộ); Điều 355, các khoản 2, 3 và 4 (tội lạm dụng chức vụ, quyền hạn chiếm đoạt tài sản); Điều 356, các khoản 2 và 3 (tội lợi dụng chức vụ, quyền hạn trong khi thi hành công vụ); Điều 357, các khoản 2 và 3 (tội lạm quyền trong khi thi hành công vụ); Điều 358, các khoản 2, 3 và 4 (tội lợi dụng chức vụ, quyền hạn gây ảnh hưởng đối với người khác để trục lợi); Điều 359, các khoản 2, 3 và 4 (tội giả mạo trong công tác); Điều 364, các khoản 2, 3 và 4 (tội đưa hối lộ); Điều 365, các khoản 2, 3 và 4 (tội làm môi giới hối lộ);\n",
+      "k) Điều 373, các khoản 3 và 4 (tội dùng nhục hình); Điều 374, các khoản 3 và 4 (tội bức cung); Điều 386, khoản 2 (tội trốn khỏi nơi giam, giữ hoặc trốn khi đang bị áp giải, dẫn giải, đang bị xét xử);\n",
+      "l) Các điều từ Điều 421 đến Điều 425 về các tội phá hoại hòa bình, chống loài người và tội phạm chiến tranh.\n",
+      "2. Phạm tội trong trường hợp lợi dụng chức vụ, quyền hạn cản trở việc phát hiện tội phạm hoặc có những hành vi khác bao che người phạm tội, thì bị phạt tù từ 02 năm đến 07 năm.\n",
+      "Điều 390. Tội không tố giác tội phạm\n",
+      "1. Người nào biết rõ một trong các tội phạm được quy định tại Điều 389 của Bộ luật này đang được chuẩn bị, đang hoặc đã được thực hiện mà không tố giác, nếu không thuộc trường hợp quy định tại khoản 2 Điều 19 của Bộ luật này, thì bị phạt cảnh cáo, phạt cải tạo không giam giữ đến 03 năm hoặc phạt tù từ 06 tháng đến 03 năm.\n",
+      "2. Người không tố giác nếu đã có hành động can ngăn người phạm tội hoặc hạn chế tác hại của tội phạm, thì có thể được miễn trách nhiệm hình sự hoặc miễn hình phạt.\n",
+      "Điều 391. Tội gây rối trật tự phiên tòa\n",
+      "1. Người nào tại phiên tòa mà thóa mạ, xúc phạm nghiêm trọng danh dự, nhân phẩm thành viên Hội đồng xét xử, những người khác có mặt tại phiên tòa hoặc có hành vi đập phá tài sản thì bị phạt tiền từ 10.000.000 đồng đến 100.000.000 đồng, phạt cải tạo không giam giữ đến 01 năm hoặc phạt tù từ 03 tháng đến 01 năm.\n",
+      "2. Phạm tội thuộc một trong các trường hợp sau đây, thì bị phạt tù từ 01 năm đến 03 năm:\n",
+      "a) Gây náo loạn phiên tòa dẫn đến phải dừng phiên tòa;\n",
+      "b) Hành hung thành viên Hội đồng xét xử.\n",
+      "--------------------------------------------------------------------------------\n",
+      "[Rank 4] index=228404, score=0.5660\n",
+      "Điều 155. Tội làm nhục người khác\n",
+      "1. Người nào xúc phạm nghiêm trọng nhân phẩm, danh dự của người khác, thì bị phạt cảnh cáo, phạt tiền từ 10.000.000 đồng đến 30.000.000 đồng hoặc phạt cải tạo không giam giữ đến 03 năm.\n",
+      "2. Phạm tội thuộc một trong các trường hợp sau đây, thì bị phạt tù từ 03 tháng đến 02 năm:\n",
+      "a) Phạm tội 02 lần trở lên;\n",
+      "b) Đối với 02 người trở lên;\n",
+      "c) Lợi dụng chức vụ, quyền hạn;\n",
+      "d) Đối với người đang thi hành công vụ;\n",
+      "đ) Đối với người dạy dỗ, nuôi dưỡng, chăm sóc, chữa bệnh cho mình;\n",
+      "e) Sử dụng mạng máy tính hoặc mạng viễn thông, phương tiện điện tử để phạm tội;\n",
+      "g) Gây rối loạn tâm thần và hành vi của nạn nhân từ 11% đến 45%.\n",
+      "3. Phạm tội thuộc một trong các trường hợp sau đây, thì bị phạt tù từ 02 năm đến 05 năm:\n",
+      "a) Gây rối loạn tâm thần và hành vi của nạn nhân 46% trở lên;\n",
+      "b) Làm nạn nhân tự sát.\n",
+      "4. Người phạm tội còn có thể bị cấm đảm nhiệm chức vụ, cấm hành nghề hoặc làm công việc nhất định từ 01 năm đến 05 năm.\n",
+      "--------------------------------------------------------------------------------\n",
+      "[Rank 5] index=143035, score=0.5470\n",
+      "Khoản 4. Người có hành vi xâm phạm thân thể, sức khỏe, tính mạng hoặc xúc phạm danh dự, nhân phẩm của người hành nghề và người khác làm việc tại cơ sở khám bệnh, chữa bệnh thì tùy theo tính chất, mức độ vi phạm mà bị xử lý vi phạm hành chính hoặc bị truy cứu trách nhiệm hình sự theo quy định của pháp luật.\n",
+      "--------------------------------------------------------------------------------\n",
+      "[Rank 6] index=57787, score=0.5443\n",
+      "\"Điều 7. Vi phạm quy định về trật tự công cộng\n",
+      "..\n",
+      "3. Phạt tiền từ 2.000.000 đồng đến 3.000.000 đồng đối với một trong những hành vi sau đây:\n",
+      "a) Có hành vi khiêu khích, trêu ghẹo, xúc phạm, lăng mạ, bôi nhọ danh dự, nhân phẩm của người khác, trừ trường hợp quy định tại điểm b khoản 2 Điều 21 và Điều 54 Nghị định này;\n",
+      "...\n",
+      "14. Biện pháp khắc phục hậu quả:\n",
+      "a) Buộc khôi phục lại tình trạng ban đầu đối với hành vi vi phạm quy định tại điểm c khoản 1, điểm l khoản 2 và điểm e khoản 4 Điều này;\n",
+      "b) Buộc cải chính thông tin sai sự thật hoặc gây nhầm lẫn đối với hành vi vi phạm quy định tại điểm a khoản 3 và điểm i khoản 4 Điều này;\n",
+      "c) Buộc xin lỗi công khai đối với hành vi vi phạm quy định tại điểm a khoản 3, các điểm d và đ khoản 5 Điều này trừ trường hợp nạn nhân có đơn không yêu cầu;\n",
+      "d) Buộc thực hiện biện pháp khắc phục tình trạng ô nhiễm môi trường đối với hành vi vi phạm quy định tại điểm h khoản 5 Điều này;\n",
+      "đ) Buộc chi trả toàn bộ chi phí khám bệnh, chữa bệnh đối với hành vi vi phạm quy định tại điểm d khoản 1, điểm c khoản 2, điểm b khoản 3 và điểm a khoản 5 Điều này.\"\n",
+      "--------------------------------------------------------------------------------\n",
+      "[Rank 7] index=57120, score=0.5337\n",
+      "Vi phạm quy định về trật tự công cộng\n",
+      "...\n",
+      "2. Phạt tiền từ 1.000.000 đồng đến 2.000.000 đồng đối với một trong những hành vi sau đây:\n",
+      "...\n",
+      "b) Tổ chức, tham gia tụ tập nhiều người ở nơi công cộng gây mất trật tự công cộng;\n",
+      "...\n",
+      "3. Phạt tiền từ 2.000.000 đồng đến 3.000.000 đồng đối với một trong những hành vi sau đây:\n",
+      "a) Có hành vi khiêu khích, trêu ghẹo, xúc phạm, lăng mạ, bôi nhọ danh dự, nhân phẩm của người khác, trừ trường hợp quy định tại điểm b khoản 2 Điều 21 và Điều 54 Nghị định này;\n",
+      "b) Tổ chức, thuê, xúi giục, lôi kéo, dụ dỗ, kích động người khác cố ý gây thương tích hoặc gây tổn hại cho sức khỏe người khác hoặc xâm phạm danh dự, nhân phẩm của người khác nhưng không bị truy cứu trách nhiệm hình sự;\n",
+      "...\n",
+      "4. Phạt tiền từ 3.000.000 đồng đến 5.000.000 đồng đối với một trong những hành vi sau đây:\n",
+      "a) Tổ chức thuê, xúi giục, lôi kéo, dụ dỗ hoặc kích động người khác gây rối, làm mất trật tự công cộng;\n",
+      "b) Mang theo trong người hoặc tàng trữ, cất giấu các loại vũ khí thô sơ, công cụ hỗ trợ hoặc các loại công cụ, phương tiện khác có khả năng sát thương; đồ vật, phương tiện giao thông nhằm mục đích gây rối trật tự công cộng, cố ý gây thương tích cho người khác;\n",
+      "...\n",
+      "5. Phạt tiền từ 5.000.000 đồng đến 8.000.000 đồng đối với một trong những hành vi sau đây:\n",
+      "a) Cố ý gây thương tích hoặc gây tổn hại cho sức khỏe của người khác nhưng không bị truy cứu trách nhiệm hình sự;\n",
+      "b) Gây rối trật tự công cộng mà có mang theo các loại vũ khí thô sơ, công cụ hỗ trợ hoặc công cụ, đồ vật, phương tiện khác có khả năng sát thương;\n",
+      "...\n",
+      "13. Hình thức xử phạt bổ sung:\n",
+      "a) Tịch thu tang vật, phương tiện vi phạm hành chính đối với hành vi vi phạm quy định tại các điểm a, d, đ và g khoản 2; điểm đ khoản 3; các đi��m b, e và i khoản 4; các điểm a, b và c khoản 5; các khoản 6 và 10 Điều này;\n",
+      "...\n",
+      "14. Biện pháp khắc phục hậu quả:\n",
+      "...\n",
+      "b) Buộc cải chính thông tin sai sự thật hoặc gây nhầm lẫn đối với hành vi vi phạm quy định tại điểm a khoản 3 và điểm i khoản 4 Điều này;\n",
+      "c) Buộc xin lỗi công khai đối với hành vi vi phạm quy định tại điểm a khoản 3, các điểm d và đ khoản 5 Điều này trừ trường hợp nạn nhân có đơn không yêu cầu;\n",
+      "...\n",
+      "đ) Buộc chi trả toàn bộ chi phí khám bệnh, chữa bệnh đối với hành vi vi phạm quy định tại điểm d khoản 1, điểm c khoản 2, điểm b khoản 3 và điểm a khoản 5 Điều này.\n",
+      "--------------------------------------------------------------------------------\n",
+      "[Rank 8] index=56183, score=0.5270\n",
+      "\"Điều 155. Tội làm nhục người khác\n",
+      "1. Người nào xúc phạm nghiêm trọng nhân phẩm, danh dự của người khác, thì bị phạt cảnh cáo, phạt tiền từ 10.000.000 đồng đến 30.000.000 đồng hoặc phạt cải tạo không giam giữ đến 03 năm.\n",
+      "2. Phạm tội thuộc một trong các trường hợp sau đây, thì bị phạt tù từ 03 tháng đến 02 năm:\n",
+      "a) Phạm tội 02 lần trở lên;\n",
+      "b) Đối với 02 người trở lên;\n",
+      "c) Lợi dụng chức vụ, quyền hạn;\n",
+      "d) Đối với người đang thi hành công vụ;\n",
+      "đ) Đối với người dạy dỗ, nuôi dưỡng, chăm sóc, chữa bệnh cho mình;\n",
+      "e) Sử dụng mạng máy tính hoặc mạng viễn thông, phương tiện điện tử để phạm tội;\n",
+      "g) Gây rối loạn tâm thần và hành vi của nạn nhân mà tỷ lệ tổn thương cơ thể từ 31% đến 60%”.\n",
+      "3. Phạm tội thuộc một trong các trường hợp sau đây, thì bị phạt tù từ 02 năm đến 05 năm:\n",
+      "a) Gây rối loạn tâm thần và hành vi của nạn nhân mà tỷ lệ tổn thương cơ thể 61% trở lên”;\n",
+      "b) Làm nạn nhân tự sát.\n",
+      "4. Người phạm tội còn có thể bị cấm đảm nhiệm chức vụ, cấm hành nghề hoặc làm công việc nhất định từ 01 năm đến 05 năm.\n",
+      "Điều 156. Tội vu khống\n",
+      "1. Người nào thực hiện một trong các hành vi sau đây, thì bị phạt tiền từ 10.000.000 đồng đến 50.000.000 đồng, phạt cải tạo không giam giữ đến 02 năm hoặc phạt tù từ 03 tháng đến 01 năm:\n",
+      "a) Bịa đặt hoặc loan truyền những điều biết rõ là sai sự thật nhằm xúc phạm nghiêm trọng nhân phẩm, danh dự hoặc gây thiệt hại đến quyền, lợi ích hợp pháp của người khác;\n",
+      "b) Bịa đặt người khác phạm tội và tố cáo họ trước cơ quan có thẩm quyền.\n",
+      "2. Phạm tội thuộc một trong các trường hợp sau đây, thì bị phạt tù từ 01 năm đến 03 năm:\n",
+      "a) Có tổ chức;\n",
+      "b) Lợi dụng chức vụ, quyền hạn;\n",
+      "c) Đối với 02 người trở lên;\n",
+      "d) Đối với ông, bà, cha, mẹ, người dạy dỗ, nuôi dưỡng, chăm sóc, giáo dục, chữa bệnh cho mình;\n",
+      "đ) Đối với người đang thi hành công vụ;\n",
+      "e) Sử dụng mạng máy tính hoặc mạng viễn thông, phương tiện điện tử để phạm tội;\n",
+      "g) Gây rối loạn tâm thần và hành vi của nạn nhân mà tỷ lệ tổn thương cơ thể từ 31% đến 60%;\n",
+      "h) Vu khống người khác phạm tội rất nghiêm trọng hoặc đặc biệt nghiêm trọng.\n",
+      "3. Phạm tội thuộc một trong các trường hợp sau đây, thì bị phạt tù từ 03 năm đến 07 năm:\n",
+      "a) Vì động cơ đê hèn;\n",
+      "b) Gây rối loạn tâm thần và hành vi của nạn nhân mà tỷ lệ tổn thương cơ thể 61% trở lên;\n",
+      "c) Làm nạn nhân tự sát.\n",
+      "4. Người phạm tội còn có thể bị phạt tiền từ 10.000.000 đồng đến 50.000.000 đồng, cấm đảm nhiệm chức vụ, cấm hành nghề hoặc làm công việc nhất định từ 01 năm đến 05 năm.\"\n",
+      "--------------------------------------------------------------------------------\n",
+      "[Rank 9] index=80022, score=0.5218\n",
+      "\"Điều 20.\n",
+      "1. Mọi người có quyền bất khả xâm phạm về thân thể, được pháp luật bảo hộ về sức khoẻ, danh dự và nhân phẩm; không bị tra tấn, bạo lực, truy bức, nhục hình hay bất kỳ hình thức đối xử nào khác xâm phạm thân thể, sức khỏe, xúc phạm danh dự, nhân phẩm.\"\n",
+      "--------------------------------------------------------------------------------\n",
+      "[Rank 10] index=52682, score=0.5203\n",
+      "\"Điều 589. Thiệt hại do tài sản bị xâm phạm\n",
+      "Thiệt hại do tài sản bị xâm phạm bao gồm:\n",
+      "1. Tài sản bị mất, bị hủy hoại hoặc bị hư hỏng.\n",
+      "2. Lợi ích gắn liền với việc sử dụng, khai thác tài sản bị mất, bị giảm sút.\n",
+      "3. Chi phí hợp lý để ngăn chặn, hạn chế và khắc phục thiệt hại.\n",
+      "4. Thiệt hại khác do luật quy định.\"\n",
+      "--------------------------------------------------------------------------------\n"
+     ]
+    }
+   ],
+   "source": [
+    "query = 'Tội xúc phạm danh dự'\n",
+    "hits  = retrieval(fine_tuned_model, query, legal_index, top_k=10)\n",
+    "\n",
+    "for h in hits:\n",
+    "    print(f\"[Rank {hits.index(h)+1}] index={h['index']}, score={h['score']:.4f}\")\n",
+    "    print(f\"{h['text']}\\n{'-'*80}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1bedd1a7",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# def search(model, query, index, k=10):\n",
+    "#     query_embedding = model.encode(\n",
+    "#         query, \n",
+    "#         convert_to_numpy=True, \n",
+    "#         normalize_embeddings=True,\n",
+    "#     ).astype(np.float32).reshape(1, -1)\n",
+    "\n",
+    "#     scores, indices = index.search(query_embedding, k*3)\n",
+    "#     hits = [{'score': scores[0][i], 'index': indices[0][i]} for i in range(len(scores[0]))]\n",
+    "#     return hits"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4ef857db",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# hits = search(\n",
+    "#     model=fine_tuned_model, \n",
+    "#     query='Hợp đồng lao động là gì?', \n",
+    "#     index=legal_index, \n",
+    "#     k=10\n",
+    "# )\n",
+    "\n",
+    "# for rank, hit in enumerate(hits):\n",
+    "#     print(f\"[Rank: {rank + 1}]\")\n",
+    "#     print(f\"(Index: {hit['index']}Score: {hit['score']:.4f})\\n\")\n",
+    "#     print(passages[hit['index']])\n",
+    "#     print('-' * 100)\n",
+    "#     print()"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "legal_doc_retrieval",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.16"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}