RAG 中 Rerank 重排序的原理、实现与工程优化

(query, document_1)
(query, document_2)
(query, document_3)
...

100000 × 1ms = 100s

20 × 1ms = 20ms

from __future__ import annotations

from dataclasses import dataclass
from typing import Iterable

import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer


@dataclass
class ScoredDocument:
    text: str
    score: float


class CrossEncoderReranker:
    def __init__(
        self,
        model_name: str = "BAAI/bge-reranker-v2-m3",
        device: str | None = None,
        max_length: int = 512,
    ) -> None:
        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
        self.max_length = max_length

        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.model = AutoModelForSequenceClassification.from_pretrained(model_name)
        self.model.to(self.device)
        self.model.eval()

    @torch.no_grad()
    def rerank(
        self,
        query: str,
        documents: list[str],
        top_k: int = 5,
        batch_size: int = 16,
        normalize_with_sigmoid: bool = True,
    ) -> list[ScoredDocument]:
        """
        对候选文档做 Cross-Encoder 重排序。

        参数:
            query: 用户问题
            documents: 向量检索召回的候选文档
            top_k: 重排序后保留的文档数量
            batch_size: 批量推理大小
            normalize_with_sigmoid: 是否把 logits 映射到 0~1

        返回:
            按相关性分数降序排列的文档列表
        """
        if not documents:
            return []

        all_scores: list[float] = []

        for batch_docs in self._batched(documents, batch_size):
            pairs = [[query, doc] for doc in batch_docs]

            inputs = self.tokenizer(
                pairs,
                padding=True,
                truncation=True,
                max_length=self.max_length,
                return_tensors="pt",
            ).to(self.device)

            logits = self.model(**inputs).logits.squeeze(-1)

            if normalize_with_sigmoid:
                scores = torch.sigmoid(logits)
            else:
                scores = logits

            all_scores.extend(scores.detach().cpu().tolist())

        scored_docs = [
            ScoredDocument(text=doc, score=float(score))
            for doc, score in zip(documents, all_scores)
        ]

        scored_docs.sort(key=lambda item: item.score, reverse=True)
        return scored_docs[:top_k]

    @staticmethod
    def _batched(items: list[str], batch_size: int) -> Iterable[list[str]]:
        for start in range(0, len(items), batch_size):
            yield items[start : start + batch_size]

query = "重疾险的等待期是多少天？"

candidate_docs = [
    "等待期内发生的疾病不在保障范围之内。",
    "等待期是保险合同的重要组成部分。",
    "等待期的计算方式从合同生效日起计算。",
    "本产品等待期为 180 天。",
]

reranker = CrossEncoderReranker()
results = reranker.rerank(query, candidate_docs, top_k=3)

for item in results:
    print(round(item.score, 4), item.text)

def filter_by_threshold(
    scored_docs: list[ScoredDocument],
    threshold: float = 0.5,
) -> list[ScoredDocument]:
    """
    丢弃低于相关性阈值的文档。
    """
    return [doc for doc in scored_docs if doc.score >= threshold]

def rag_retrieve(
    query: str,
    vector_store,
    reranker: CrossEncoderReranker,
    top_k_recall: int = 20,
    top_k_rerank: int = 5,
    threshold: float = 0.5,
) -> list[ScoredDocument]:
    """
    RAG 检索链路：
    向量召回 -> Cross-Encoder 重排序 -> 阈值过滤。
    """
    recall_results = vector_store.similarity_search(query, k=top_k_recall)
    candidate_texts = [doc.page_content for doc in recall_results]

    reranked_docs = reranker.rerank(
        query=query,
        documents=candidate_texts,
        top_k=top_k_rerank,
    )

    final_docs = filter_by_threshold(
        scored_docs=reranked_docs,
        threshold=threshold,
    )

    return final_docs

def build_answer(query: str, final_docs: list[ScoredDocument], llm) -> str:
    if not final_docs:
        return "知识库中没有找到足够相关的内容，无法基于现有资料回答该问题。"

    context = "\n\n".join(doc.text for doc in final_docs)

    prompt = f"""
请只基于给定资料回答问题。
如果资料不足以回答，请说明资料中没有相关信息。

资料：
{context}

问题：
{query}
"""

    return llm.invoke(prompt)

{
  "query": "重疾险的等待期是多少天？",
  "doc": "本产品等待期为 180 天。",
  "label": 1
}

from sklearn.metrics import precision_score, recall_score, f1_score


def search_best_threshold(
    labels: list[int],
    scores: list[float],
    start: float = 0.3,
    end: float = 0.8,
    step: float = 0.05,
) -> dict:
    """
    遍历多个阈值，选择 F1 分数最高的阈值。
    """
    best = {
        "threshold": None,
        "precision": 0.0,
        "recall": 0.0,
        "f1": 0.0,
    }

    threshold = start
    while threshold <= end + 1e-9:
        preds = [1 if score >= threshold else 0 for score in scores]

        precision = precision_score(labels, preds, zero_division=0)
        recall = recall_score(labels, preds, zero_division=0)
        f1 = f1_score(labels, preds, zero_division=0)

        if f1 > best["f1"]:
            best = {
                "threshold": round(threshold, 4),
                "precision": round(precision, 4),
                "recall": round(recall, 4),
                "f1": round(f1, 4),
            }

        threshold += step

    return best

(query, positive_doc, negative_doc)

import json
from pathlib import Path


def contains_answer(doc_text: str, answer: str) -> bool:
    """
    简单启发式判断：文档中是否包含答案片段。
    生产环境可以替换成规则匹配、人工复核或语义判断。
    """
    normalized_doc = doc_text.replace(" ", "")
    normalized_answer = answer.replace(" ", "")

    return normalized_answer in normalized_doc


def prepare_reranker_finetune_data(
    qa_pairs: list[dict],
    vector_store,
    output_path: str,
    recall_k: int = 50,
) -> None:
    """
    生成 Reranker 微调所需的三元组数据。

    qa_pairs 中每条数据格式示例：
    {
        "query": "重疾险等待期是多少天？",
        "answer": "180 天",
        "positive_doc": "本产品等待期为 180 天。"
    }
    """
    examples: list[dict] = []

    for qa in qa_pairs:
        query = qa["query"]
        answer = qa["answer"]
        positive_doc = qa["positive_doc"]

        candidates = vector_store.similarity_search(query, k=recall_k)

        hard_negatives: list[str] = []

        for candidate in candidates:
            candidate_text = candidate.page_content

            if candidate_text == positive_doc:
                continue

            if not contains_answer(candidate_text, answer):
                hard_negatives.append(candidate_text)

        if not hard_negatives:
            continue

        hardest_negative = hard_negatives[0]

        examples.append(
            {
                "query": query,
                "pos": [positive_doc],
                "neg": [hardest_negative],
            }
        )

    output_file = Path(output_path)
    output_file.parent.mkdir(parents=True, exist_ok=True)

    with output_file.open("w", encoding="utf-8") as f:
        for example in examples:
            f.write(json.dumps(example, ensure_ascii=False) + "\n")

    print(f"生成 {len(examples)} 条微调样本：{output_file}")

{"query": "重疾险等待期是多少天？", "pos": ["本产品等待期为 180 天。"], "neg": ["等待期内发生的疾病不在保障范围之内。"]}
{"query": "轻症赔付比例是多少？", "pos": ["轻度恶性肿瘤按基本保额的 20% 给付。"], "neg": ["轻症疾病需符合合同约定定义。"]}

results = reranker.rerank(
    query=query,
    documents=candidate_docs,
    top_k=5,
    batch_size=16,
)

def pack_context(
    docs: list[ScoredDocument],
    max_chars: int = 4000,
) -> str:
    chunks: list[str] = []
    total = 0

    for doc in docs:
        text = doc.text.strip()

        if total + len(text) > max_chars:
            break

        chunks.append(text)
        total += len(text)

    return "\n\n".join(chunks)

Bi-Encoder 向量召回 Top-20
        ↓
Cross-Encoder Rerank
        ↓
取 Top-5
        ↓
阈值过滤
        ↓
送入 LLM