Day 37: Reranking Cho Production RAG

(query, candidate_chunk) -> relevance_score

Retriever: tìm rộng để không bỏ sót.
Reranker: chấm đắt hơn nhưng chính xác hơn trên tập candidate nhỏ.
Context builder: lấy vài chunk tốt nhất, giữ citation và giới hạn token.

query -> embedding_q
doc   -> embedding_d
score = similarity(embedding_q, embedding_d)

[query, doc] -> transformer -> relevance_score

Title: Chính sách nghỉ phép 2026
Section: HR > Leave Policy > Annual Leave
Version: 2026-01
Text: Nhân viên full-time có 12 ngày nghỉ phép năm...

from sentence_transformers import CrossEncoder

model = CrossEncoder("BAAI/bge-reranker-v2-m3", max_length=512)

query = "Nhân viên full-time được nghỉ phép bao nhiêu ngày trong năm 2026?"
passages = [
    "Title: Chính sách nghỉ phép 2026\nText: Nhân viên full-time có 12 ngày nghỉ phép năm.",
    "Title: Quy định bảo mật\nText: Nhân viên phải đổi mật khẩu mỗi 90 ngày.",
]

scores = model.predict([(query, passage) for passage in passages])

import os

import cohere

co = cohere.Client(token=os.environ["COHERE_API_KEY"])

response = co.v2.rerank(
    model="rerank-v4.0-pro",
    query="Nhân viên full-time được nghỉ phép bao nhiêu ngày?",
    documents=[
        "Title: Chính sách nghỉ phép\nText: Nhân viên full-time có 12 ngày nghỉ phép năm.",
        "Title: Chính sách bảo mật\nText: Không chia sẻ mật khẩu cho người khác.",
    ],
    top_n=2,
)

for item in response.results:
    print(item.index, item.relevance_score)

user query
  -> normalize query
  -> BM25 top 50 + vector top 50
  -> merge bằng RRF
  -> filter tenant/ACL/deleted/index_version
  -> dedupe theo document_id/chunk_id/text_hash
  -> truncate text theo max tokens per doc
  -> rerank top 50 hoặc top 100
  -> lấy top 5-10
  -> context builder giữ citation
  -> LLM answer

{
  "query_id": "q001",
  "query": "Nhân viên full-time có bao nhiêu ngày nghỉ phép năm 2026?",
  "relevant_chunk_ids": ["hr_leave_2026:chunk_003"]
}

from __future__ import annotations

import logging
import os
import time
from dataclasses import dataclass
from typing import Callable, Protocol, Sequence

logger = logging.getLogger(__name__)


@dataclass(frozen=True)
class AccessContext:
    tenant_id: str
    roles: frozenset[str]


@dataclass(frozen=True)
class CandidateChunk:
    id: str
    document_id: str
    chunk_id: str
    tenant_id: str
    acl_roles: frozenset[str]
    title: str
    section_path: tuple[str, ...]
    text: str
    source_uri: str
    page_start: int | None
    page_end: int | None
    retrieval_score: float
    retrieval_rank: int
    retriever: str

    def text_for_reranker(self, max_chars: int = 2_500) -> str:
        section = " > ".join(self.section_path)
        body = self.text[:max_chars]
        return (
            f"Title: {self.title}\n"
            f"Section: {section}\n"
            f"Source: {self.source_uri}\n"
            f"Text: {body}"
        )


@dataclass(frozen=True)
class RankedChunk:
    chunk: CandidateChunk
    rerank_score: float
    final_rank: int


class Reranker(Protocol):
    name: str
    model_version: str

    def score(self, query: str, candidates: Sequence[CandidateChunk]) -> list[float]:
        """Return one relevance score per candidate, same order as input."""


class BgeCrossEncoderReranker:
    name = "bge-cross-encoder"

    def __init__(
        self,
        model_name: str = "BAAI/bge-reranker-v2-m3",
        max_length: int = 512,
        device: str | None = None,
    ) -> None:
        from sentence_transformers import CrossEncoder

        self.model_version = model_name
        self.model = CrossEncoder(model_name, max_length=max_length, device=device)

    def score(self, query: str, candidates: Sequence[CandidateChunk]) -> list[float]:
        pairs = [(query, candidate.text_for_reranker()) for candidate in candidates]
        scores = self.model.predict(pairs)
        return [float(score) for score in scores]


class CohereReranker:
    name = "cohere-rerank"

    def __init__(
        self,
        api_key: str | None = None,
        model: str = "rerank-v4.0-pro",
        max_tokens_per_doc: int = 1_200,
    ) -> None:
        import cohere

        self.model_version = model
        self.client = cohere.Client(token=api_key or os.environ["COHERE_API_KEY"])
        self.max_tokens_per_doc = max_tokens_per_doc

    def score(self, query: str, candidates: Sequence[CandidateChunk]) -> list[float]:
        documents = [candidate.text_for_reranker() for candidate in candidates]
        response = self.client.v2.rerank(
            model=self.model_version,
            query=query,
            documents=documents,
            top_n=len(documents),
            max_tokens_per_doc=self.max_tokens_per_doc,
        )

        scores = [float("-inf")] * len(candidates)
        for result in response.results:
            scores[result.index] = float(result.relevance_score)
        return scores


HybridRetrieveFn = Callable[[str, int], list[CandidateChunk]]


def has_access(candidate: CandidateChunk, access: AccessContext) -> bool:
    if candidate.tenant_id != access.tenant_id:
        return False
    return bool(candidate.acl_roles & access.roles)


def dedupe_candidates(candidates: Sequence[CandidateChunk]) -> list[CandidateChunk]:
    best_by_chunk: dict[str, CandidateChunk] = {}
    for candidate in candidates:
        previous = best_by_chunk.get(candidate.id)
        if previous is None or candidate.retrieval_score > previous.retrieval_score:
            best_by_chunk[candidate.id] = candidate

    return sorted(
        best_by_chunk.values(),
        key=lambda item: (item.retrieval_rank, -item.retrieval_score),
    )


def fallback_rank(candidates: Sequence[CandidateChunk], final_k: int) -> list[RankedChunk]:
    return [
        RankedChunk(chunk=candidate, rerank_score=float("nan"), final_rank=index + 1)
        for index, candidate in enumerate(candidates[:final_k])
    ]


def retrieve_then_rerank(
    *,
    query: str,
    access: AccessContext,
    hybrid_retrieve: HybridRetrieveFn,
    reranker: Reranker,
    retrieve_k: int = 100,
    rerank_k: int = 50,
    final_k: int = 8,
    rerank_timeout_ms: int = 800,
) -> list[RankedChunk]:
    started = time.perf_counter()

    raw_candidates = hybrid_retrieve(query, retrieve_k)
    permitted = [item for item in raw_candidates if has_access(item, access)]
    deduped = dedupe_candidates(permitted)
    rerank_input = deduped[:rerank_k]

    if not rerank_input:
        return []

    try:
        scores = reranker.score(query, rerank_input)
    except Exception:
        logger.exception(
            "reranker_failed",
            extra={
                "reranker": reranker.name,
                "model_version": reranker.model_version,
                "retrieve_k": retrieve_k,
                "rerank_k": rerank_k,
            },
        )
        return fallback_rank(rerank_input, final_k)

    elapsed_ms = (time.perf_counter() - started) * 1000
    if elapsed_ms > rerank_timeout_ms:
        logger.warning(
            "reranker_slow",
            extra={
                "elapsed_ms": round(elapsed_ms, 2),
                "timeout_ms": rerank_timeout_ms,
                "reranker": reranker.name,
            },
        )
        return fallback_rank(rerank_input, final_k)

    ranked_pairs = sorted(
        zip(rerank_input, scores, strict=True),
        key=lambda item: item[1],
        reverse=True,
    )

    return [
        RankedChunk(chunk=chunk, rerank_score=float(score), final_rank=index + 1)
        for index, (chunk, score) in enumerate(ranked_pairs[:final_k])
    ]

from collections.abc import Iterable


def recall_at_k(ranked_ids: Sequence[str], relevant_ids: set[str], k: int) -> float:
    if not relevant_ids:
        return 0.0
    retrieved = set(ranked_ids[:k])
    return len(retrieved & relevant_ids) / len(relevant_ids)


def mrr_at_k(ranked_ids: Sequence[str], relevant_ids: set[str], k: int) -> float:
    for rank, chunk_id in enumerate(ranked_ids[:k], start=1):
        if chunk_id in relevant_ids:
            return 1.0 / rank
    return 0.0


def evaluate_ranker(
    run: dict[str, list[str]],
    qrels: dict[str, set[str]],
    recall_k: int = 5,
    mrr_k: int = 10,
) -> dict[str, float]:
    recall_scores = []
    mrr_scores = []

    for query_id, relevant_ids in qrels.items():
        ranked_ids = run.get(query_id, [])
        recall_scores.append(recall_at_k(ranked_ids, relevant_ids, recall_k))
        mrr_scores.append(mrr_at_k(ranked_ids, relevant_ids, mrr_k))

    total = max(len(qrels), 1)
    return {
        f"Recall@{recall_k}": sum(recall_scores) / total,
        f"MRR@{mrr_k}": sum(mrr_scores) / total,
    }

Hybrid retrieval = recall layer
Reranker         = precision layer
Context builder  = token/citation layer
LLM              = answer layer

| Pipeline | Recall@5 | MRR@10 | nDCG@10 | p95 search ms | p95 rerank ms | p95 total ms | Cost/1K | Note |
|---|---:|---:|---:|---:|---:|---:|---:|---|
| Hybrid only | | | | | 0 | | | |
| Hybrid + BGE rerank top 50 | | | | | | | | |
| Hybrid + Cohere rerank top 50 | | | | | | | | |
| Hybrid + rerank top 100 | | | | | | | | |

BM25 top 50 + Vector top 50
  -> RRF merge
  -> ACL filter
  -> dedupe
  -> rerank top 50
  -> final top 5/10
  -> evaluate Recall@k / MRR

pip install sentence-transformers numpy

pip install cohere
export COHERE_API_KEY="..."

QUERIES = [
    {
        "query_id": "q001",
        "query": "Nhân viên full-time có bao nhiêu ngày nghỉ phép năm 2026?",
        "tags": ["policy", "version", "vietnamese"],
    },
    {
        "query_id": "q002",
        "query": "VPN error 809 xử lý như thế nào?",
        "tags": ["keyword-heavy", "incident"],
    },
    {
        "query_id": "q003",
        "query": "Can contractor access production database?",
        "tags": ["english", "acl", "security"],
    },
]

QRELS = {
    "q001": {"hr_leave_2026:chunk_003"},
    "q002": {"it_vpn_runbook:chunk_007"},
    "q003": {"security_access_policy:chunk_011"},
}

def run_hybrid_only(query: str, top_k: int = 10) -> list[str]:
    bm25_hits = bm25_search(query, top_k=50)
    vector_hits = vector_search(query, top_k=50)
    merged = reciprocal_rank_fusion([bm25_hits, vector_hits])
    permitted = acl_filter(merged)
    deduped = dedupe(permitted)
    return [hit.chunk_id for hit in deduped[:top_k]]

from sentence_transformers import CrossEncoder


class LocalBgeReranker:
    def __init__(self, model_name: str = "BAAI/bge-reranker-v2-m3") -> None:
        self.model = CrossEncoder(model_name, max_length=512)

    def rerank(self, query: str, candidates: list[dict], top_n: int = 10) -> list[dict]:
        pairs = [
            (
                query,
                "Title: {title}\nSection: {section}\nText: {text}".format(
                    title=item.get("title", ""),
                    section=" > ".join(item.get("section_path", [])),
                    text=item["text"][:2500],
                ),
            )
            for item in candidates
        ]
        scores = self.model.predict(pairs)

        ranked = sorted(
            zip(candidates, scores, strict=True),
            key=lambda item: float(item[1]),
            reverse=True,
        )
        return [
            {
                **candidate,
                "rerank_score": float(score),
                "final_rank": index + 1,
            }
            for index, (candidate, score) in enumerate(ranked[:top_n])
        ]

def run_hybrid_plus_rerank(query: str, reranker: LocalBgeReranker) -> list[str]:
    bm25_hits = bm25_search(query, top_k=50)
    vector_hits = vector_search(query, top_k=50)
    merged = reciprocal_rank_fusion([bm25_hits, vector_hits])
    permitted = acl_filter(merged)
    deduped = dedupe(permitted)

    candidates = [hit.to_dict() for hit in deduped[:50]]
    ranked = reranker.rerank(query, candidates, top_n=10)
    return [item["chunk_id"] for item in ranked]

import os

import cohere


class CohereRerankClient:
    def __init__(self, model: str = "rerank-v4.0-pro") -> None:
        self.client = cohere.Client(token=os.environ["COHERE_API_KEY"])
        self.model = model

    def rerank(self, query: str, candidates: list[dict], top_n: int = 10) -> list[dict]:
        documents = [
            "Title: {title}\nText: {text}".format(
                title=item.get("title", ""),
                text=item["text"][:2500],
            )
            for item in candidates
        ]
        response = self.client.v2.rerank(
            model=self.model,
            query=query,
            documents=documents,
            top_n=top_n,
            max_tokens_per_doc=1200,
        )

        ranked = []
        for rank, result in enumerate(response.results, start=1):
            item = dict(candidates[result.index])
            item["rerank_score"] = float(result.relevance_score)
            item["final_rank"] = rank
            ranked.append(item)
        return ranked

def recall_at_k(ranked_ids: list[str], relevant_ids: set[str], k: int) -> float:
    if not relevant_ids:
        return 0.0
    return len(set(ranked_ids[:k]) & relevant_ids) / len(relevant_ids)


def mrr_at_k(ranked_ids: list[str], relevant_ids: set[str], k: int) -> float:
    for rank, chunk_id in enumerate(ranked_ids[:k], start=1):
        if chunk_id in relevant_ids:
            return 1.0 / rank
    return 0.0


def evaluate(run: dict[str, list[str]], qrels: dict[str, set[str]]) -> dict[str, float]:
    recall_5 = []
    mrr_10 = []
    for query_id, relevant_ids in qrels.items():
        ranked_ids = run.get(query_id, [])
        recall_5.append(recall_at_k(ranked_ids, relevant_ids, 5))
        mrr_10.append(mrr_at_k(ranked_ids, relevant_ids, 10))

    total = len(qrels) or 1
    return {
        "Recall@5": sum(recall_5) / total,
        "MRR@10": sum(mrr_10) / total,
    }

hybrid_run = {
    item["query_id"]: run_hybrid_only(item["query"], top_k=10)
    for item in QUERIES
}

reranker = LocalBgeReranker()
rerank_run = {
    item["query_id"]: run_hybrid_plus_rerank(item["query"], reranker)
    for item in QUERIES
}

print("Hybrid only:", evaluate(hybrid_run, QRELS))
print("Hybrid + rerank:", evaluate(rerank_run, QRELS))

import statistics
import time


def benchmark(fn, queries: list[dict], repeat: int = 3) -> dict[str, float]:
    latencies = []
    for _ in range(repeat):
        for item in queries:
            started = time.perf_counter()
            fn(item["query"])
            latencies.append((time.perf_counter() - started) * 1000)

    latencies = sorted(latencies)
    p95_index = int(0.95 * (len(latencies) - 1))
    return {
        "avg_ms": statistics.mean(latencies),
        "p50_ms": statistics.median(latencies),
        "p95_ms": latencies[p95_index],
        "max_ms": max(latencies),
    }

| Pipeline | Recall@5 | MRR@10 | p50 ms | p95 ms | Cost/1K query | Ghi chú |
|---|---:|---:|---:|---:|---:|---|
| Hybrid only | | | | | | |
| Hybrid + BGE top 20 | | | | | | |
| Hybrid + BGE top 50 | | | | | | |
| Hybrid + Cohere top 50 | | | | | | |

Có thể dùng reranker trong production cho nhóm query policy/support vì Hybrid + BGE top 50 tăng MRR@10 từ 0.61 lên 0.78 và Recall@5 từ 0.70 lên 0.84 trên 120 query eval. p95 retrieval tăng từ 420ms lên 780ms, vẫn dưới SLA 1.200ms. Hệ thống đã filter tenant/ACL trước rerank, có timeout 800ms và fallback về RRF rank. Giai đoạn đầu bật bằng feature flag cho 20% internal traffic, log before/after rank và theo dõi citation correctness.

Nên dùng reranker vì nó thông minh hơn vector search.