hoailebads/Qwen3-Reranker-0.6B-VLSP-Legal-LoRA

Qwen3-Reranker-0.6B — VLSP Vietnamese Legal Reranker (LoRA)

LoRA adapter cho Qwen/Qwen3-Reranker-0.6B, fine-tune để xếp hạng lại (rerank) điều luật tiếng Việt cho bài toán truy hồi văn bản pháp luật VLSP: cho một câu hỏi pháp lý và một danh sách điều luật ứng viên, model chấm điểm mức liên quan của từng cặp (câu hỏi, điều luật).

Đây là checkpoint tốt nhất (checkpoint-600) sau khi quét toàn bộ checkpoint của lần train v6 trên bộ dữ liệu 23 hard-negative/query — cân bằng tốt nhất giữa R@1 và R@10.

Bản này là phiên bản 0.6B nhẹ: cùng pipeline và cùng format input với bản hoailebads/Qwen3-Reranker-8B-VLSP-Legal-LoRA nhưng nhỏ hơn ~13 lần, chạy được trên GPU phổ thông. Đổi lại độ chính xác thấp hơn bản 8B.

• 📦 Code, pipeline, script tune & tạo submission: https://github.com/hoaileba/Qwen-Retrieval-Tuning
• 🗂️ Dữ liệu train + eval: https://huggingface.co/datasets/hoailebads/VLSP-Legal-Reranker-Train
• 📄 Tài liệu kỹ thuật chi tiết: docs/PIPELINE.md

---

Điểm quan trọng: đây là đường SEQ_CLS, KHÔNG phải yes/no

Qwen3-Reranker-0.6B gốc chấm điểm theo đường generative (logit("yes") − logit("no")). Adapter này không dùng đường đó. Model được nạp qua AutoModelForSequenceClassification(num_labels=1), điểm là logit vô hướng ở token cuối, và score head tuyến tính (Linear(hidden → 1, bias=False)) được huấn luyện cùng LoRA và lưu trong file adapter dưới key base_model.model.score.weight (PEFT modules_to_save).

Hệ quả — muốn điểm đúng thì bắt buộc:

1. Nạp base bằng AutoModelForSequenceClassification(num_labels=1, ignore_mismatched_sizes=True) → tạo score head mới (random), rồi PeftModel.from_pretrained(...) sẽ ghi đè score head bằng trọng số đã train (vì nó nằm trong modules_to_save). Nạp lên AutoModelForCausalLM hoặc dùng prompt yes/no sẽ ra điểm vô nghĩa.
2. Đặt pad_token_id và padding_side="left": SEQ_CLS pool hidden state ở token cuối cùng không phải pad. Thiếu bước này (pad sai / pad bên phải) → pool nhầm token → điểm sai hoàn toàn.

Code dưới đây khớp đúng với script inference trong repo (reranker/eval_rerank_ckpt_nofaiss.py).

---

Kết quả

Tập eval: 219 câu hỏi VLSP (ground-truth dedup theo aid), rerank top-100 ứng viên do bi-encoder Qwen3-Embedding-0.6B + dual-LoRA sinh ra. Đơn vị: %.

Cấu hình	R@1	R@3	R@5	R@10
Retrieval thuần (không rerank)	45.36	63.51	69.06	81.13
Base Qwen3-Reranker-0.6B (chưa fine-tune, đường yes/no)	40.49	60.81	68.00	78.50
+ LoRA này (checkpoint-600)	55.18	71.39	78.23	85.12

• Rerank thêm +9.82 R@1 và +3.99 R@10 so với retrieval thuần.
• Base 0.6B (đường yes/no) chưa fine-tune còn thấp hơn cả retrieval → fine-tune là bắt buộc.
• Quét toàn bộ checkpoint: checkpoint-600 cho R@10 cao nhất (85.12) và R@1 gần cao nhất. Bảng đầy đủ: eval_checkpoint_sweep.csv.

So với bản 8B

Trên private leaderboard VLSP, đổi reranker 0.6B → 8B đáng giá đáng kể; bản 8B đạt F2MACRO 0.7360 (vượt SOTA VLSP 0.7261). Bản 0.6B này là lựa chọn nhẹ/chi phí thấp: nếu cần độ chính xác tối đa hãy dùng bản 8B; nếu cần tốc độ/VRAM thấp, dùng bản này.

---

Cách dùng

import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from peft import PeftModel

BASE    = "Qwen/Qwen3-Reranker-0.6B"
ADAPTER = "hoailebads/Qwen3-Reranker-0.6B-VLSP-Legal-LoRA"
MAX_DOC_LEN = 1024
INSTRUCTION = ("Given a Vietnamese legal question, retrieve the most relevant "
               "legal article that directly answers the question")

# ── Tokenizer: BẮT BUỘC left-pad + có pad_token (SEQ_CLS pool token cuối) ──
tok = AutoTokenizer.from_pretrained(BASE, trust_remote_code=True)
if tok.pad_token is None:
    tok.pad_token = tok.eos_token
tok.padding_side = "left"        # pooling lấy token cuối
tok.truncation_side = "left"     # giữ <eos> ở cuối chuỗi

# ── Model: SEQ_CLS(num_labels=1) → PeftModel nạp LoRA + score head ──
try:
    import flash_attn  # noqa
    attn = "flash_attention_2"
except ImportError:
    attn = "sdpa"

base = AutoModelForSequenceClassification.from_pretrained(
    BASE, num_labels=1, torch_dtype=torch.bfloat16, trust_remote_code=True,
    attn_implementation=attn,
    ignore_mismatched_sizes=True,        # score head mới ≠ lm_head gốc
)
base.config.pad_token_id = tok.pad_token_id   # BẮT BUỘC — pool đúng token cuối
model = PeftModel.from_pretrained(base, ADAPTER)   # ghi đè score head bằng trọng số đã train
model.eval().cuda()

# (tuỳ chọn) kiểm tra score head đã nạp — không phải random:
#   from safetensors.torch import load_file
#   sd = load_file(hf_hub_download(ADAPTER, "adapter_model.safetensors"))
#   assert "base_model.model.score.weight" in sd     # trọng số score head có trong adapter

def build_input(query: str, doc: str) -> str:
    # Document pre-truncate TRƯỚC khi ghép chuỗi (khớp đúng lúc train / eval)
    ids = tok(doc, add_special_tokens=False)["input_ids"]
    if len(ids) > MAX_DOC_LEN:
        doc = tok.decode(ids[:MAX_DOC_LEN], skip_special_tokens=True)
    return (f"Instruct: {INSTRUCTION}\n"
            f"query: {query}\n"
            f"document: {doc}") + tok.eos_token

@torch.no_grad()
def score(query: str, docs: list[str]) -> list[float]:
    texts = [build_input(query, d) for d in docs]
    enc = tok(texts, add_special_tokens=False, padding=True, truncation=True,
              max_length=MAX_DOC_LEN + 200, return_tensors="pt").to(model.device)
    with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
        logits = model(input_ids=enc["input_ids"],
                       attention_mask=enc["attention_mask"]).logits   # [B, 1]
    return logits.squeeze(-1).float().cpu().tolist()

query = "Phạm nhân không biết chữ có được tạo điều kiện học văn hóa để xóa mù chữ không?"
docs = ["Điều 31. Chế độ học tập, học nghề của phạm nhân ...", "Điều 5. Nguyên tắc ..."]
ranked = sorted(zip(docs, score(query, docs)), key=lambda x: -x[1])

⚠️ Điểm là logit THÔ — chỉ so được TRONG cùng một query

score head xuất logit chưa hiệu chuẩn. So sánh điểm giữa các query khác nhau là sai. Đừng dùng ngưỡng tuyệt đối để quyết định "có liên quan hay không" — hãy dùng hình dạng điểm trong từng query (margin hoặc softmax) để chọn số đáp án động. Code sẵn: submission/make_dynamic_submission.py và submission/make_softmax_submission.py trong repo GitHub.

---

Chi tiết huấn luyện

Base model	Qwen/Qwen3-Reranker-0.6B (decoder-only)
Kiểu	AutoModelForSequenceClassification, num_labels=1, task_type=SEQ_CLS
LoRA	r=32, alpha=64, dropout=0.1
Target modules	q_proj, k_proj, v_proj, o_proj, gate_proj, up_proj, down_proj
modules_to_save	score (đầu điểm tuyến tính — huấn luyện đầy đủ, lưu trong adapter)
Loss	listwise cross-entropy, positive ở index 0 (InfoNCE, KHÔNG in-batch negative)
Hard negative	23/query → 24 passage/query, mine bằng dense retriever
Max length	doc pre-truncate 1024 token, tổng 1024 + 200
Batch	batch_size=5 × grad_accum=2
Optimizer	lr 1e-4, weight decay 0.05, warmup ratio 0.06, max_grad_norm=0.5
Precision	bf16 + gradient checkpointing
Checkpoint chọn	checkpoint-600 (step 600, epoch ≈ 1.2; train loss 6.73 → 0.13)
Seed	42

Config đầy đủ: train_config.json.

Nội dung repo này

File
adapter_model.safetensors	LoRA weights + score head (base_model.model.score.weight, 77 MB)
adapter_config.json	config PEFT (base_model_name_or_path = Qwen/Qwen3-Reranker-0.6B)
train_config.json	toàn bộ hyperparameter của lần train v6 (0.6B)
eval_checkpoint_sweep.csv	R@k của từng checkpoint + baseline
reranker_meta.json	metadata (input format, instruction, score head)

Trạng thái optimizer/scheduler không được upload — chỉ cần cho việc resume training.

Giới hạn

• Chỉ dành cho truy hồi điều luật tiếng Việt; miền khác cần re-tune.
• Trần recall của pha retrieval là R@100 = 95.21 trên eval → reranker không thể cứu ~5% câu hỏi không có đáp án đúng trong 100 ứng viên.
• Bản 0.6B kém hơn bản 8B về độ chính xác — chọn theo ràng buộc tốc độ/VRAM.
• Instruction và format input phải khớp chính xác như trên, sai lệch sẽ giảm điểm.

License

Adapter phát hành theo Apache-2.0. Trọng số base Qwen/Qwen3-Reranker-0.6B theo license riêng của Qwen.