Update new

README.md

@@ -1,19 +1,3 @@
----
-license: mit
-datasets:
-- another-symato/VMTEB-Zalo-legel-retrieval-wseg
-language:
-- vi
-base_model:
-- bkai-foundation-models/vietnamese-bi-encoder
-pipeline_tag: feature-extraction
-tags:
-- feature-extraction
-- embedding
-- legal
-- law
-- vietnamese
----
 # TELEN: Temporal Evolving Legal Embedding Network
 
 > **Vietnamese legal text embedding with meta-learning for continuous adaptation to new laws.**
@@ -66,15 +50,21 @@ L2-Normalized Embedding [768-dim]
 | Model | NDCG@3 | NDCG@5 | NDCG@10 | MRR@3 | MRR@5 | MRR@10 |
 |---|---|---|---|---|---|---|
 | **BM25** (lexical) | 0.6753 | 0.7173 | 0.7250 | 0.6683 | 0.6928 | 0.6990 |
-| **PhoBERT-base-v2** (dense) | 0.5866 | 0.6360 | 0.6505 | 0.5657 | 0.5970 | 0.6059 |
-| **DEk21** (dense) | 0.7900 | 0.8127 | 0.8344 | 0.7660 | 0.7785 | 0.7865 |
-| **TELEN** (dense) | 0.9036 | 0.9138 | 0.9132 | 0.8830 | 0.8878 | 0.8878 |
-| **TELEN + CE re-rank** (dense) | **0.9346** | **0.9339** | **0.9238** | **0.9199** | **0.9223** | **0.9223** |
+| **PhoBERT-base-v2** (monolingual dense) | 0.5866 | 0.6360 | 0.6505 | 0.5657 | 0.5970 | 0.6059 |
+| **multilingual-E5-base** (multilingual dense) | 0.4675 | 0.4888 | 0.5157 | 0.4327 | 0.4452 | 0.4573 |
+| **BAAI/bge-m3** (multilingual dense, 1024d) | 0.4668 | 0.5129 | 0.5452 | 0.4407 | 0.4657 | 0.4802 |
+| **DEk21** (legal dense) | 0.7900 | 0.8127 | 0.8344 | 0.7660 | 0.7785 | 0.7865 |
+| **TELEN** (adaptive dense) | 0.9036 | 0.9138 | 0.9132 | 0.8830 | 0.8878 | 0.8878 |
+| **TELEN + CE re-rank** (adaptive dense) | **0.9346** | **0.9339** | **0.9238** | **0.9199** | **0.9223** | **0.9223** |
+
+> **Key insight:** Multilingual SOTA models (multilingual-E5, BGE-M3) score **below even BM25** on Vietnamese legal text, confirming that domain and language specialization trumps generic multilingual pre-training for legal retrieval.
 
 ### Relative Improvement
 
 | Baseline | NDCG@3 | NDCG@5 | NDCG@10 | MRR@10 |
 |---|---|---|---|---|
+| vs multilingual-E5 | **+93.3%** | **+86.9%** | **+77.1%** | **+94.1%** |
+| vs BGE-M3 | **+93.6%** | **+78.2%** | **+67.5%** | **+84.9%** |
 | vs PhoBERT | **+59.3%** | **+46.8%** | **+42.0%** | **+52.2%** |
 | vs DEk21 | **+18.3%** | **+14.9%** | **+10.7%** | **+17.3%** |
 
@@ -124,7 +114,7 @@ python train_ce.py
 ### Evaluation
 
 ```bash
-# Full benchmark (TELEN vs BM25/PhoBERT/DEk21)
+# Full benchmark (TELEN vs BM25/PhoBERT/mE5/BGE-M3/DEk21)
 python eval.py
 
 # TELEN + Cross-encoder re-ranking (MRR-optimized)
@@ -232,4 +222,4 @@ MIT License — see [LICENSE](LICENSE) file for details.
 ## Acknowledgments
 
 - `bkai-foundation-models/vietnamese-bi-encoder` — backbone bi-encoder
-- `huyydangg/DEk21_hcmute_embedding` — baseline comparison - `vinai/phobert-base-v2` — used in cross-encoder re-ranker
+- `huyydangg/DEk21_hcmute_embedding` — baseline comparison - `vinai/phobert-base-v2` — used in cross-encoder re-ranker

eval.py

@@ -6,6 +6,8 @@ Metrics: NDCG@3, NDCG@5, NDCG@10, MRR@3, MRR@5, MRR@10
 Baselines:
   - BM25 (lexical retrieval)
   - Frozen PhoBERT (vinai/phobert-base-v2)
+  - multilingual-E5-base (intfloat/multilingual-e5-base)
+  - BGE-M3 (BAAI/bge-m3)
   - DEk21 (huyydangg/DEk21_hcmute_embedding)
   - TELEN (ours)
 
@@ -19,6 +21,7 @@ import random, numpy as np, torch, torch.nn.functional as F
 from tqdm import tqdm
 from collections import defaultdict
 from sentence_transformers import SentenceTransformer
+from transformers import AutoModel, AutoTokenizer
 from pyvi import ViTokenizer
 
 from src.telern.config import TELENConfig
@@ -35,13 +38,15 @@ config = TELENConfig()
 def wseg(text):
     return ViTokenizer.tokenize(text.replace("_", " "))
 
-def evaluate_model(name, encode_fn, queries, corpus, corpus_ids, corpus_law_ids):
+def evaluate_model(name, encode_fn, queries, corpus, corpus_ids, corpus_law_ids, corpus_encode_fn=None):
     """Generic evaluation for any embedding model."""
+    if corpus_encode_fn is None:
+        corpus_encode_fn = encode_fn
     print(f"\n  [{name}] Encoding corpus ({len(corpus)} docs)...")
     c_embs = []
     for i in range(0, len(corpus), 64):
         batch = [d["text"] for d in corpus[i:i+64]]
-        embs = encode_fn(batch)
+        embs = corpus_encode_fn(batch)
         if isinstance(embs, np.ndarray): embs = torch.tensor(embs)
         c_embs.append(embs.cpu())
     c_embs = torch.cat(c_embs, dim=0)
@@ -73,21 +78,59 @@ print(f"Test: {len(queries)} queries, {len(corpus)} docs, {test_df['law_id'].nun
 results = {}
 
 # ── BM25 ──
-print("\n[1/4] BM25")
+print("\n[1/6] BM25")
 results["BM25"] = evaluate_bm25(queries, corpus)
 
 # ── PhoBERT ──
-print("\n[2/4] Frozen PhoBERT")
+print("\n[2/6] Frozen PhoBERT")
 phobert = FrozenPhoBERT()
 results["PhoBERT"] = evaluate_model("PhoBERT", lambda texts: phobert.encode(texts, batch_size=64), queries, corpus, corpus_ids, corpus_law_ids)
 
 # ── DEk21 ──
-print("\n[3/4] DEk21 (SOTA)")
+print("\n[3/6] DEk21 (legal SOTA)")
 dek21 = SentenceTransformer("huyydangg/DEk21_hcmute_embedding", device=device)
 results["DEk21"] = evaluate_model("DEk21", lambda texts: dek21.encode([wseg(t) for t in texts], batch_size=64, show_progress_bar=False, normalize_embeddings=True, convert_to_tensor=True), queries, corpus, corpus_ids, corpus_law_ids)
 
+# ── multilingual-E5-base ──
+print("\n[4/6] multilingual-E5-base")
+e5_tokenizer = AutoTokenizer.from_pretrained("intfloat/multilingual-e5-base")
+e5_model = AutoModel.from_pretrained("intfloat/multilingual-e5-base").to(device)
+e5_model.eval()
+def e5_encode(texts, prefix="query: "):
+    prefixed = [prefix + t for t in texts]
+    enc = e5_tokenizer(prefixed, padding=True, truncation=True, max_length=512, return_tensors="pt")
+    with torch.no_grad():
+        hidden = e5_model(input_ids=enc["input_ids"].to(device), attention_mask=enc["attention_mask"].to(device)).last_hidden_state
+        mask = enc["attention_mask"].unsqueeze(-1).float().to(device)
+        pooled = (hidden * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1e-9)
+    return F.normalize(pooled, p=2, dim=1)
+results["multilingual-e5"] = evaluate_model("mE5",
+    lambda texts: e5_encode(texts),  # queries: "query: " prefix
+    queries, corpus, corpus_ids, corpus_law_ids,
+    corpus_encode_fn=lambda texts: e5_encode(texts, prefix="passage: "))
+del e5_model, e5_tokenizer; torch.cuda.empty_cache()
+
+# ── BGE-M3 ──
+print("\n[5/6] BAAI/bge-m3")
+bge_tokenizer = AutoTokenizer.from_pretrained("BAAI/bge-m3")
+bge_model = AutoModel.from_pretrained("BAAI/bge-m3").to(device)
+bge_model.eval()
+def bge_encode(texts, add_prefix=True):
+    if add_prefix:
+        texts = ["Represent this sentence for searching relevant passages: " + t for t in texts]
+    enc = bge_tokenizer(texts, padding=True, truncation=True, max_length=512, return_tensors="pt")
+    with torch.no_grad():
+        hidden = bge_model(input_ids=enc["input_ids"].to(device), attention_mask=enc["attention_mask"].to(device)).last_hidden_state
+        cls_emb = hidden[:, 0, :]
+    return F.normalize(cls_emb, p=2, dim=1)
+results["bge-m3"] = evaluate_model("BGE-M3",
+    lambda texts: bge_encode(texts, add_prefix=True),  # queries: with instruction
+    queries, corpus, corpus_ids, corpus_law_ids,
+    corpus_encode_fn=lambda texts: bge_encode(texts, add_prefix=False))  # passages: no prefix
+del bge_model, bge_tokenizer; torch.cuda.empty_cache()
+
 # ── TELEN ──
-print("\n[4/4] TELEN (Ours)")
+print("\n[6/6] TELEN (Ours)")
 telen = create_model(config).to(device)
 ckpt = torch.load(config.output_dir + "/telen_best.pt", map_location=device, weights_only=False)
 telen.hypernetwork.load_state_dict(ckpt["hypernetwork"])
@@ -109,14 +152,16 @@ print("=" * 75)
 h = f"{'Method':<15}"
 for m in [3,5,10]: h += f" {'NDCG@'+str(m):>10} {'MRR@'+str(m):>10}"
 print(h); print("-"*len(h))
-for name in ["BM25", "PhoBERT", "DEk21", "TELEN"]:
-    r = f"{name:<15}"
+for name in ["BM25", "PhoBERT", "multilingual-e5", "bge-m3", "DEk21", "TELEN"]:
+    display = {"multilingual-e5": "mE5-base", "bge-m3": "BGE-M3"}.get(name, name)
+    r = f"{display:<15}"
     for m in [3,5,10]: r += f" {results[name][f'ndcg@{m}']:>10.4f} {results[name][f'mrr@{m}']:>10.4f}"
     print(r)
 
 print("\n--- Relative Improvement over Baselines ---")
-for baseline in ["PhoBERT", "DEk21"]:
-    print(f"  TELEN vs {baseline}:")
+for baseline in ["PhoBERT", "multilingual-e5", "bge-m3", "DEk21"]:
+    display = {"multilingual-e5": "mE5-base", "bge-m3": "BGE-M3"}.get(baseline, baseline)
+    print(f"  TELEN vs {display}:")
     for m in [3,5,10]:
         ni = (results["TELEN"][f"ndcg@{m}"] / max(results[baseline][f"ndcg@{m}"], 1e-6) - 1) * 100
         mi = (results["TELEN"][f"mrr@{m}"] / max(results[baseline][f"mrr@{m}"], 1e-6) - 1) * 100