10.RAG 检索与重排序

实现余弦相似度计算函数，并基于暴力检索完成向量检索：给定查询向量和文档向量矩阵，返回相似度最高的 top‑k 个文档的索引及分数。¶

1.1 余弦相似度原理¶

余弦相似度通过计算两个向量夹角的余弦值来衡量它们的方向一致性，其值域为[-1,1]，值越大表示越相似。对于文档检索，通常将查询和文档都表示为高维稀疏向量（如TF‑IDF）或稠密向量（如BERT嵌入）。余弦相似度公式为：

如果向量已经L2归一化，则余弦相似度退化为内积。

1.2 暴力检索的实现¶

暴力检索（Brute-Force Search）直接计算查询向量与所有文档向量之间的相似度，然后降序排列取前k个。其时间复杂度为O(N·D)，其中N为文档数，D为向量维度。当N在十万以内且D不高时，暴力检索完全可以接受，且精确度100%。

1.3 代码实现¶

import numpy as np

def cosine_similarity(vec1: np.ndarray, vec2: np.ndarray) -> float:
    """计算两个向量的余弦相似度"""
    dot = np.dot(vec1, vec2)
    norm1 = np.linalg.norm(vec1)
    norm2 = np.linalg.norm(vec2)
    if norm1 == 0 or norm2 == 0:
        return 0.0
    return dot / (norm1 * norm2)

def brute_force_search(query_vec: np.ndarray, doc_matrix: np.ndarray,
                       top_k: int = 10, normalize: bool = False):
    """
    暴力检索 top-k 最相似文档
    Args:
        query_vec: (D,) 查询向量
        doc_matrix: (N, D) 文档向量矩阵，每行一篇文档
        top_k: 返回的文档数
        normalize: 是否提前对文档矩阵做L2归一化以加速（内积代替余弦）
    Returns:
        indices: (k,) 相似度最高的文档索引
        scores: (k,) 对应的余弦相似度分数
    """
    if normalize:
        # 如果doc_matrix已经归一化，则余弦相似度等于内积
        # 同时query_vec也需要归一化
        query_vec = query_vec / np.linalg.norm(query_vec)
        scores = np.dot(doc_matrix, query_vec)   # (N,)
    else:
        # 逐文档计算余弦相似度（也可以向量化）
        # 向量化写法：doc_norms = np.linalg.norm(doc_matrix, axis=1)
        # scores = np.dot(doc_matrix, query_vec) / (doc_norms * np.linalg.norm(query_vec))
        # 但更高效的是上面的公式，这里用逐次计算演示
        scores = np.array([cosine_similarity(query_vec, doc_vec) for doc_vec in doc_matrix])

    # 获取top-k索引（argsort默认升序，取最后k个再反转即为降序）
    if top_k >= len(scores):
        top_k = len(scores)
    # 使用argpartition可以更高效地只取topk，但这里用简单排序
    top_indices = np.argsort(scores)[::-1][:top_k]
    top_scores = scores[top_indices]
    return top_indices, top_scores

# 示例
if __name__ == "__main__":
    np.random.seed(42)
    N, D = 1000, 128
    doc_matrix = np.random.randn(N, D)        # 随机文档向量
    query = np.random.randn(D)
    indices, scores = brute_force_search(query, doc_matrix, top_k=5, normalize=True)
    print("Top-5 indices:", indices)
    print("Top-5 scores:", scores)

1.4 优化与注意事项¶

数据归一化：事先对文档向量进行L2归一化，查询时也归一化，则相似度退化为内积，速度更快。
批量检索：使用矩阵乘法（doc_matrix @ query_vec）可以利用BLAS加速，一次计算所有相似度。
内存占用：文档矩阵通常较大，若无法全部放入内存，可考虑分块计算。
精度：暴力检索返回的是真实Top‑K，不存在近似误差，适用于召回阶段或小规模数据。

构建一个简单的倒排索引：输入文档集合（文档 ID 和分词后的词列表），输出每个词对应的文档 ID 列表及词频。¶

2.1 倒排索引原理¶

倒排索引是搜索引擎的核心数据结构。对于每个词项（term），记录包含它的文档ID列表（posting list）以及在该文档中的词频（TF）。通常还会记录位置信息用于短语查询。这里我们构建一个基础版本，包含文档ID和词频。

2.2 构建过程¶

遍历每个文档，对文档中的每个词进行统计。使用字典 inverted_index，键为词项，值为另一个字典 {doc_id: tf}。

2.3 代码实现¶

from collections import defaultdict
from typing import List, Dict, Tuple

def build_inverted_index(documents: List[Tuple[int, List[str]]]) -> Dict[str, Dict[int, int]]:
    """
    构建倒排索引
    Args:
        documents: 列表，每个元素为 (doc_id, tokens) ，tokens是分词后的词列表
    Returns:
        倒排索引: {term: {doc_id: term_freq}}
    """
    inverted_index = defaultdict(lambda: defaultdict(int))
    for doc_id, tokens in documents:
        # 统计本文档各词频率
        tf_map = defaultdict(int)
        for token in tokens:
            tf_map[token] += 1
        # 写入倒排索引
        for token, tf in tf_map.items():
            inverted_index[token][doc_id] = tf
    return dict(inverted_index)

# 示例
docs = [
    (1, ["hello", "world", "hello"]),
    (2, ["hello", "python"]),
    (3, ["world", "world", "python", "code"])
]
index = build_inverted_index(docs)
for term, posting in index.items():
    print(f"'{term}': {posting}")

输出：

'hello': {1: 2, 2: 1}
'world': {1: 1, 3: 2}
'python': {2: 1, 3: 1}
'code': {3: 1}

2.4 扩展¶

可以同时存储文档总词数，用于后续BM25计算。
为支持BM25，还需要记录文档长度（词数），可以单独保存一个字典 doc_lengths。
实际系统中倒排索引通常以压缩格式存储，如使用VarInt编码压缩倒排列表。

实现 BM25 评分函数：输入查询词项、文档词频、文档长度、平均文档长度、参数 k1 和 b，计算该文档对查询的 BM25 分数。¶

3.1 BM25 公式¶

BM25 是概率检索模型的经典实现，其公式为：

其中 N 为总文档数，nt 为包含 t 的文档数。

3.2 代码实现¶

import math

def compute_idf(N: int, nt: int) -> float:
    """计算IDF，nt为包含该词的文档数"""
    if nt == 0:
        return 0.0
    return math.log((N - nt + 0.5) / (nt + 0.5))

def bm25_score_single(tf: int, doc_len: int, avgdl: float,
                      k1: float = 1.5, b: float = 0.75) -> float:
    """计算单个查询词对单个文档的BM25贡献（不含IDF）"""
    numerator = tf * (k1 + 1)
    denominator = tf + k1 * (1 - b + b * (doc_len / avgdl))
    return numerator / denominator

def bm25_score(query_terms: list, tf_map: dict, doc_len: int, avgdl: float,
               idf_dict: dict, k1=1.5, b=0.75) -> float:
    """
    计算一个文档对查询的BM25分数
    Args:
        query_terms: 查询的词项列表
        tf_map: 该文档的词频字典 {term: tf}
        doc_len: 文档长度
        avgdl: 平均文档长度
        idf_dict: 每个词的IDF值字典
    """
    score = 0.0
    for t in query_terms:
        tf = tf_map.get(t, 0)
        if tf == 0:
            continue
        idf = idf_dict.get(t, 0.0)
        if idf == 0.0:
            continue
        score += idf * bm25_score_single(tf, doc_len, avgdl, k1, b)
    return score

3.3 示例¶

假设有3篇文档，查询"hello world"。

docs = [
    (1, ["hello", "world", "hello"]),
    (2, ["hello", "python"]),
    (3, ["world", "world", "python"])
]
# 计算文档长度
doc_lengths = {1:3, 2:2, 3:3}
avgdl = sum(doc_lengths.values())/len(doc_lengths)   # 2.666
N = len(docs)
# 统计nt
nt_hello = sum(1 for _, tokens in docs if "hello" in tokens)  # 2
nt_world = 2
idf_hello = compute_idf(N, nt_hello)
idf_world = compute_idf(N, nt_world)
idf_dict = {"hello": idf_hello, "world": idf_world}

# 文档1的tf map
tf1 = {"hello":2, "world":1}
score = bm25_score(["hello","world"], tf1, doc_lengths[1], avgdl, idf_dict)
print("Doc1 BM25:", score)

基于已构建的倒排索引，实现完整的 BM25 检索：接受查询字符串，计算所有可能相关文档的 BM25 分数，返回 top‑k 结果。¶

4.1 整体流程¶

对查询进行分词，得到查询词项列表。
从倒排索引中取出每个查询词对应的 posting list（文档ID和词频），得到候选文档集合。
获取每个候选文档的长度，以及全局平均文档长度、文档总数、每个词的IDF等信息。
对每个候选文档，遍历查询词，利用已存储的tf计算BM25分数。
将所有候选文档的分数降序排列，返回top‑k。

4.2 代码实现¶

class BM25Retriever:
    def __init__(self, k1=1.5, b=0.75):
        self.k1 = k1
        self.b = b
        self.inverted_index = {}   # {term: {doc_id: tf}}
        self.doc_lengths = {}      # {doc_id: length}
        self.N = 0                 # 文档总数
        self.avgdl = 0.0
        self.idf_cache = {}        # {term: idf}

    def index_documents(self, documents: List[Tuple[int, List[str]]]):
        """构建倒排索引并计算统计量"""
        total_len = 0
        self.N = len(documents)
        for doc_id, tokens in documents:
            # 统计词频
            tf_map = defaultdict(int)
            for t in tokens:
                tf_map[t] += 1
            # 写入索引
            for t, tf in tf_map.items():
                if t not in self.inverted_index:
                    self.inverted_index[t] = {}
                self.inverted_index[t][doc_id] = tf
            # 记录长度
            self.doc_lengths[doc_id] = len(tokens)
            total_len += len(tokens)
        self.avgdl = total_len / self.N if self.N > 0 else 0.0
        # 预计算IDF
        for term, postings in self.inverted_index.items():
            nt = len(postings)
            self.idf_cache[term] = compute_idf(self.N, nt)

    def search(self, query: str, top_k: int = 10) -> List[Tuple[int, float]]:
        """检索 top-k 文档"""
        # 分词（简单空格分割）
        query_terms = query.strip().split()
        # 收集所有相关文档及其每个查询词的tf
        candidate_docs = defaultdict(lambda: defaultdict(int))  # {doc_id: {term: tf}}
        for term in query_terms:
            if term in self.inverted_index:
                for doc_id, tf in self.inverted_index[term].items():
                    candidate_docs[doc_id][term] = tf
        # 计算每个候选文档的BM25分数
        scores = {}
        for doc_id, term_tf_map in candidate_docs.items():
            doc_len = self.doc_lengths[doc_id]
            score = 0.0
            for term in query_terms:
                tf = term_tf_map.get(term, 0)
                if tf == 0:
                    continue
                idf = self.idf_cache.get(term, 0.0)
                score += idf * self._bm25_tf(tf, doc_len)
            scores[doc_id] = score
        # 排序
        sorted_docs = sorted(scores.items(), key=lambda x: x[1], reverse=True)[:top_k]
        return sorted_docs

    def _bm25_tf(self, tf, doc_len):
        return tf * (self.k1 + 1) / (tf + self.k1 * (1 - self.b + self.b * (doc_len / self.avgdl)))

4.3 使用示例¶

retriever = BM25Retriever()
retriever.index_documents(docs)
results = retriever.search("hello world", top_k=2)
print(results)

4.4 性能优化¶

实际实现中，可以在检索时动态计算IDF（需要N和nt），预计算可节省时间。
候选文档可能很多，可采用WAND等动态剪枝算法提前终止低分文档的计算。
倒排索引中存储文档长度，避免检索时查表。

实现倒数排名融合（Reciprocal Rank Fusion, RRF）算法：输入多个检索器的排序列表（文档 ID 排名序列），参数 k，输出合并后的最终排序列表。¶

5.1 RRF 原理¶

RRF 是一种简单而有效的排名融合方法，不依赖各检索器分数的尺度。其公式为：

5.2 实现¶

from typing import List, Dict

def reciprocal_rank_fusion(ranked_lists: List[List[int]], k: float = 60.0) -> List[Tuple[int, float]]:
    """
    RRF 融合多个排序列表
    Args:
        ranked_lists: 列表的列表，每个内层列表为按排名顺序的文档ID序列
        k: RRF参数
    Returns:
        按RRF分数降序的 (doc_id, score) 列表
    """
    rrf_scores: Dict[int, float] = defaultdict(float)
    for ranked_list in ranked_lists:
        for rank, doc_id in enumerate(ranked_list, start=1):
            rrf_scores[doc_id] += 1.0 / (k + rank)
    # 按分数降序排序
    sorted_docs = sorted(rrf_scores.items(), key=lambda x: x[1], reverse=True)
    return sorted_docs

# 示例
list1 = [1, 3, 5, 7]   # 检索器1的排名
list2 = [2, 1, 3, 6]   # 检索器2的排名
list3 = [5, 1, 2, 4]
merged = reciprocal_rank_fusion([list1, list2, list3], k=60)
print("RRF merged:", merged)

5.3 分析¶

RRF 不受各检索器分数范围影响，因此特别适用于融合稠密检索（如向量相似度0~1）和稀疏检索（如BM25分数范围不固定）的结果。
常数 kk 用于控制高排名文档的权重：kk 越小，高排名文档的贡献越大。通常设为60是经验值。
对于未出现在某个检索器列表中的文档，其排名视为无穷大，贡献为0。

实现基于 Min‑Max 归一化的线性加权分数融合：给定稠密检索分数和稀疏检索分数两列，分别归一化后按权重 α 求和排序。¶

6.1 原理¶

当两种检索器返回的分数范围不一致时（例如向量相似度在0.7~0.9，BM25在5~20），直接加权会导致一种分数主导。Min‑Max归一化将分数线性映射到[0,1]区间：

6.2 代码实现¶

def minmax_normalize(scores: List[float]) -> List[float]:
    """Min-Max归一化，返回[0,1]区间"""
    if not scores:
        return []
    min_s = min(scores)
    max_s = max(scores)
    if max_s == min_s:
        return [0.5] * len(scores)   # 所有分数相同
    return [(s - min_s) / (max_s - min_s) for s in scores]

def linear_weighted_fusion(dense_scores: Dict[int, float],
                           sparse_scores: Dict[int, float],
                           alpha: float = 0.5) -> List[Tuple[int, float]]:
    """
    融合两个检索器的分数（文档ID为键的字典），返回按融合分降序列表
    alpha: 稠密检索的权重
    """
    # 收集所有文档ID
    all_doc_ids = set(dense_scores.keys()) | set(sparse_scores.keys())
    # 提取分数列表（保持顺序以便归一化后映射）
    doc_ids = list(all_doc_ids)
    raw_dense = [dense_scores.get(doc, 0.0) for doc in doc_ids]
    raw_sparse = [sparse_scores.get(doc, 0.0) for doc in doc_ids]

    # 分别归一化
    norm_dense = minmax_normalize(raw_dense)
    norm_sparse = minmax_normalize(raw_sparse)

    # 加权求和
    fused_scores = {}
    for i, doc_id in enumerate(doc_ids):
        fused = alpha * norm_dense[i] + (1 - alpha) * norm_sparse[i]
        fused_scores[doc_id] = fused

    # 降序排序
    return sorted(fused_scores.items(), key=lambda x: x[1], reverse=True)

# 示例
dense = {1: 0.9, 2: 0.8, 3: 0.6}
sparse = {1: 12.5, 2: 8.0, 4: 15.0}
result = linear_weighted_fusion(dense, sparse, alpha=0.6)
print(result)

6.3 注意事项¶

归一化是针对“当前查询”的候选文档集进行的，因此不同查询之间的融合分数不可直接比较。
如果某个检索器未返回某文档，通常赋予一个缺省分数（如0或该检索器的最小分数），但这样可能不公平。更好的做法是只考虑同时出现在两个列表中的文档，或为缺失值赋予一个较低的归一化值（如0）。
参数 α 可以通过验证集调整。

编写一个 Reranker 类，封装交叉编码器模型：输入查询和候选文档列表，使用一个预训练的交叉编码器逐对打分，返回按分数降序排列的文档列表。¶

7.1 交叉编码器原理¶

交叉编码器将查询和文档拼接后输入Transformer，通过自注意力进行深度交互，输出一个相关性分数（通常是一个标量，经过线性层+softmax得到）。相比双塔模型（分别编码后内积），交叉编码器精度高但速度慢，适合对粗筛后的少量候选进行重排序。

常用模型：cross-encoder/ms-marco-MiniLM-L-6-v2、BAAI/bge-reranker-base等。

7.2 实现¶

from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
from typing import List, Tuple

class Reranker:
    def __init__(self, model_name: str = "cross-encoder/ms-marco-MiniLM-L-6-v2", device: str = None):
        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.model = AutoModelForSequenceClassification.from_pretrained(model_name).to(self.device)
        self.model.eval()

    def rerank(self, query: str, documents: List[str],
               batch_size: int = 32, max_length: int = 512) -> List[Tuple[str, float]]:
        """
        对候选文档重排序
        Args:
            query: 查询字符串
            documents: 候选文档内容列表
            batch_size: 批量大小
            max_length: 截断长度
        Returns:
            (document, score) 降序列表
        """
        scores = []
        # 构造输入对
        pairs = [(query, doc) for doc in documents]
        # 分批编码
        for i in range(0, len(pairs), batch_size):
            batch = pairs[i:i+batch_size]
            # 分词
            inputs = self.tokenizer(batch, padding=True, truncation=True,
                                    max_length=max_length, return_tensors="pt")
            inputs = {k: v.to(self.device) for k, v in inputs.items()}
            with torch.no_grad():
                logits = self.model(**inputs).logits.squeeze(-1)  # (batch,)
                # 如果是二分类，logits是原始分数，可直接作为相关性分数；也可softmax
                # 这里直接使用logits
                scores.extend(logits.cpu().tolist())
        # 排序
        doc_score_pairs = list(zip(documents, scores))
        doc_score_pairs.sort(key=lambda x: x[1], reverse=True)
        return doc_score_pairs

# 示例
reranker = Reranker()
docs = [
    "The capital of France is Paris.",
    "France is a country in Europe.",
    "Paris is known for the Eiffel Tower."
]
results = reranker.rerank("What is the capital of France?", docs)
for doc, score in results:
    print(f"{score:.4f} {doc}")

7.3 注意事项¶

交叉编码器的输入长度限制通常为512 tokens，超长文档需要截断。对于长文档，可先分段用MaxP策略（见下一题）或只取开头部分。
分数可以是原始logits，也可softmax转成概率，但排序只需相对大小。
批量推理可以显著提升速度。

针对长文档实现 MaxP 段落聚合策略：将每个文档分成多个段落，每个段落与查询打分（可用假设的相似度函数），取最大段落分数作为文档分数。¶

8.1 原理¶

对于长文档（如学术论文、法律合同），直接将其整体与查询匹配可能会稀释相关性。MaxP（Maximum Passage）策略将文档分割为多个段落（passage），每个段落独立与查询计算相关性分数，然后取所有段落中的最高分作为该文档的最终得分。这保证了只要文档中有一个段落高度相关，文档就能被检索到。

8.2 实现¶

这里使用一个简单的相似度函数（如TF‑IDF余弦相似度）来模拟段落打分，实际应用中可替换为交叉编码器或稠密检索。

import numpy as np
from typing import List, Dict, Tuple
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity

def maxp_retrieve(query: str, documents: Dict[int, str],
                  segment_size: int = 200, overlap: int = 50) -> List[Tuple[int, float]]:
    """
    MaxP段落聚合检索
    Args:
        query: 查询字符串
        documents: {doc_id: full_text} 字典
        segment_size: 每个段落的字符数（或词语数）
        overlap: 段落之间重叠字符数
    Returns:
        按分数降序的 (doc_id, max_score) 列表
    """
    # 1. 生成所有文档的段落，并记录每个段落所属文档
    passages = []
    passage_doc_map = []   # 段落索引 -> doc_id
    for doc_id, text in documents.items():
        # 简单按固定长度切分（可改用滑动窗口或按句子）
        start = 0
        while start < len(text):
            end = start + segment_size
            passage = text[start:end]
            passages.append(passage)
            passage_doc_map.append(doc_id)
            if end >= len(text):
                break
            start = end - overlap
            if start < 0:
                start = 0

    # 2. 使用TF‑IDF向量化所有段落和查询
    #    （此处为了演示，实际可用预训练模型）
    vectorizer = TfidfVectorizer().fit(passages + [query])
    passage_vectors = vectorizer.transform(passages)
    query_vector = vectorizer.transform([query])

    # 3. 计算每个段落与查询的余弦相似度
    scores = cosine_similarity(query_vector, passage_vectors).flatten()

    # 4. 聚合：每个文档取最大段落分数
    doc_max_scores: Dict[int, float] = defaultdict(float)
    for idx, doc_id in enumerate(passage_doc_map):
        score = scores[idx]
        if score > doc_max_scores[doc_id]:
            doc_max_scores[doc_id] = score

    # 5. 排序返回
    sorted_docs = sorted(doc_max_scores.items(), key=lambda x: x[1], reverse=True)
    return sorted_docs

# 示例
docs = {
    1: "The quick brown fox jumps over the lazy dog. " * 20,   # 长文档
    2: "A fox is a wild animal. Dogs are domesticated pets.",
    3: "Python programming is fun. " * 50
}
results = maxp_retrieve("fox dog", docs)
print("MaxP results:", results)

8.3 扩展与优化¶

段落切分方式可以基于句子、滑动窗口、或使用段落分割模型。
打分函数可以替换为更强大的Cross‑Encoder或双塔模型。
MaxP策略特别适用于RAG系统，能够在文档级索引中定位最相关的片段。
除了MaxP，还可以使用MeanP（平均分）、SumP等聚合方法，但MaxP能更好地捕捉局部强相关。

实现 ColBERT 风格的晚交互 MaxSim 计算¶

原理

ColBERT 是一种“晚交互”检索模型。与将查询和文档分别编码为单一向量并直接计算相似度的双塔模型不同，ColBERT 为查询的每个 token 生成一个向量，为文档的每个 token 也生成一个向量。检索时，对于查询的每个 token 向量，在文档所有 token 向量中寻找余弦相似度最大的那个，然后将这些最大值求和，得到该文档的最终相关性分数。这个过程称为 MaxSim。

计算公式如下：

这种方法保留了 token 级别的细粒度匹配信号，同时在离线阶段可以将文档 token 向量提前计算并存储，检索时只需与查询 token 向量进行轻量级的 MaxSim 计算。

实现

import torch
import torch.nn.functional as F

def cosine_similarity_matrix(query_vectors, doc_vectors):
    """
    计算两组向量之间的余弦相似度矩阵。
    query_vectors: [Lq, D] 查询token向量
    doc_vectors: [Ld, D] 文档token向量
    返回: [Lq, Ld] 相似度矩阵
    """
    # 归一化向量
    q_norm = F.normalize(query_vectors, p=2, dim=-1)
    d_norm = F.normalize(doc_vectors, p=2, dim=-1)
    # 矩阵乘法得到余弦相似度
    return torch.mm(q_norm, d_norm.t())

def colbert_maxsim(query_vectors, doc_vectors):
    """
    Colbert 晚交互 MaxSim 分数。
    query_vectors: (Lq, D) tensor
    doc_vectors: (Ld, D) tensor
    返回: float 标量分数
    """
    sim_matrix = cosine_similarity_matrix(query_vectors, doc_vectors)  # [Lq, Ld]
    # 对每个查询token取文档token中的最大值
    max_per_query_token, _ = torch.max(sim_matrix, dim=1)  # [Lq]
    score = torch.sum(max_per_query_token).item()
    return score

# 示例
if __name__ == "__main__":
    Lq, Ld, D = 5, 10, 128
    Q = torch.randn(Lq, D)   # 5个查询token向量
    D_mat = torch.randn(Ld, D)  # 10个文档token向量
    score = colbert_maxsim(Q, D_mat)
    print(f"MaxSim score: {score:.4f}")

说明

实际 ColBERT 还会对查询和文档向量进行长度归一化（L2归一化）后计算余弦相似度，这里已包含。
查询 token 通常包含特殊标记（如 [CLS]、[Q] 等），文档 token 也类似，但这不影响计算。
该 MaxSim 分数不直接是概率，通常在检索时直接用于排序。在训练时可能会叠加温度系数或线性层。
在离线检索阶段，文档 token 向量库可以预先构建（如存入 FAISS），在线查询时只需计算 MaxSim，无需模型推理。

实现一个混合检索器，同时执行稠密向量检索和 BM25 检索，并使用 RRF 合并两者的结果。¶

原理

混合检索结合了稠密检索的语义理解能力和稀疏检索（BM25）的精确关键词匹配优势。实现思路：预先构建稠密向量索引（例如使用 FAISS 或简单的暴力检索）和 BM25 倒排索引。对于用户查询，分别从两个检索器获取各自的排序列表（文档 ID 及其原始分数或排名）。然后使用倒数排名融合（RRF）算法将两个列表合并为一个最终排序。RRF 公式：

其中 rank 从 1 开始，未出现的文档该项贡献为 0。

实现

import numpy as np
from typing import List, Dict, Tuple
from collections import defaultdict
# 复用前面第1题的暴力检索函数 brute_force_search
# 和第2题的倒排索引构建以及第3/4题的BM25检索器（BM25Retriever）

class HybridRetriever:
    def __init__(self, dense_vectors: np.ndarray, bm25_retriever, k_rrf=60):
        """
        dense_vectors: (N, D) 文档稠密向量矩阵
        bm25_retriever: BM25Retriever 实例（已经索引文档）
        k_rrf: RRF 参数
        """
        self.dense_vectors = dense_vectors
        self.bm25 = bm25_retriever
        self.k_rrf = k_rrf

    def search(self, query_str: str, query_dense_vec: np.ndarray,
               top_k_dense: int = 100, top_k_sparse: int = 100,
               final_top_k: int = 10) -> List[Tuple[int, float]]:
        """
        混合检索
        Args:
            query_str: 查询字符串（用于BM25）
            query_dense_vec: (D,) 查询的稠密向量
            top_k_dense: 稠密检索返回的候选数
            top_k_sparse: 稀疏检索返回的候选数
            final_top_k: 最终返回的文档数
        Returns:
            按RRF分数降序的 (doc_id, rrf_score) 列表
        """
        # 1. 稠密检索（暴力）
        dense_indices, dense_scores = brute_force_search(
            query_dense_vec, self.dense_vectors, top_k=top_k_dense, normalize=True
        )
        # 转换为排序列表（文档ID序列，按排名顺序）
        dense_ranked = [int(idx) for idx in dense_indices]

        # 2. 稀疏检索 BM25
        sparse_results = self.bm25.search(query_str, top_k=top_k_sparse)  # 返回 [(doc_id, score)]
        sparse_ranked = [doc_id for doc_id, _ in sparse_results]

        # 3. RRF 融合
        rrf_scores = defaultdict(float)
        # 处理稠密排名
        for rank, doc_id in enumerate(dense_ranked, start=1):
            rrf_scores[doc_id] += 1.0 / (self.k_rrf + rank)
        # 处理稀疏排名
        for rank, doc_id in enumerate(sparse_ranked, start=1):
            rrf_scores[doc_id] += 1.0 / (self.k_rrf + rank)

        # 按RRF分数降序排序
        sorted_docs = sorted(rrf_scores.items(), key=lambda x: x[1], reverse=True)
        return sorted_docs[:final_top_k]

示例代码框架

# 假设已有文档集合、稠密向量矩阵、BM25检索器
# 查询
query_text = "hello world"
query_vec = np.random.randn(128)   # 实际应使用embedding模型生成
results = hybrid_retriever.search(query_text, query_vec)
print(results)

注意事项

稠密检索通常采用近似最近邻（ANN）索引（如FAISS）来加速，这里使用了暴力检索作为演示。
BM25检索器需实现第4题中的 BM25Retriever 类，其 search 方法返回 [(doc_id, bm25_score)]。
RRF 的 k 参数常用 60，可调整。
为了提升效率，可以将稠密索引和BM25索引分开构建，并在每次查询时并行执行。

对检索结果列表进行去重，保留文档 ID 第一次出现的位置，同时可将分数保留为多次出现的最高分。¶

原理

在融合多个检索源或进行多次检索时，同一文档可能在结果列表中出现多次。去重规则通常为：按文档 ID 去重，保留排名最靠前的那一条（即第一次出现的位置），并且可以记录该文档在各次出现中的最高分数。实现时遍历列表，用字典记录已见过的ID，如果重复则跳过并更新最高分（可选）。

实现

from typing import List, Tuple

def dedup_results(results: List[Tuple[int, float]], keep_highest_score: bool = False):
    """
    对结果列表去重，保持原始顺序（保留第一次出现的ID）。
    results: [(doc_id, score), ...]
    keep_highest_score: 如果为True，则当重复出现时，保留最高分数（但不改变顺序位置）
    """
    seen = set()
    unique_results = []
    score_map = {}  # 记录每个ID的当前最高分
    for doc_id, score in results:
        if doc_id not in seen:
            seen.add(doc_id)
            unique_results.append((doc_id, score))
            if keep_highest_score:
                score_map[doc_id] = score
        else:
            # 如果允许保留最高分，且当前分更高，则更新已存在的记录的分数
            if keep_highest_score and score > score_map.get(doc_id, float('-inf')):
                # 找到该doc_id在unique_results中的位置并更新分数
                for i, (uid, _) in enumerate(unique_results):
                    if uid == doc_id:
                        unique_results[i] = (doc_id, score)
                        score_map[doc_id] = score
                        break
    return unique_results

# 示例
res = [(1, 0.9), (2, 0.8), (1, 0.95), (3, 0.7), (2, 0.6)]
deduped = dedup_results(res, keep_highest_score=True)
print(deduped)  # [(1, 0.95), (2, 0.8), (3, 0.7)]

说明

默认情况下保留第一次出现的位置和该位置的分数，这在 RRF 或分数融合后排序的列表中是合适的，因为第一次出现往往排名更高。
如果需要保留最高分并原地更新已存在的分数，上述代码提供了实现。
该函数也适用于普通排序列表的去重。

计算 MRR@k：给定真实相关文档 ID 集合和系统返回的 top‑k 文档 ID 列表，输出倒数排名均值。¶

原理

MRR（Mean Reciprocal Rank）衡量的是第一个相关文档在结果列表中的排名。对于单个查询，倒数排名 RR = 1 / rank，其中 rank 是第一个相关文档出现的位置（从1开始）；如果没有相关文档，RR = 0。MRR 是所有查询 RR 的平均值。MRR@k 表示只考虑前 k 个结果。

实现

def mrr_at_k(relevant_ids: set, retrieved_ids: List[int], k: int = None) -> float:
    """
    计算单个查询的倒数排名。
    Args:
        relevant_ids: 该查询的真实相关文档ID集合
        retrieved_ids: 系统返回的文档ID列表（按排名顺序）
        k: 截断值，默认全部
    Returns:
        RR值 (float)
    """
    if k is not None:
        retrieved_ids = retrieved_ids[:k]
    for rank, doc_id in enumerate(retrieved_ids, start=1):
        if doc_id in relevant_ids:
            return 1.0 / rank
    return 0.0

def mean_reciprocal_rank(queries_relevant: List[set], queries_retrieved: List[List[int]], k: int = 10) -> float:
    """计算多个查询的MRR@k"""
    total_rr = 0.0
    for rel_set, ret_list in zip(queries_relevant, queries_retrieved):
        total_rr += mrr_at_k(rel_set, ret_list, k)
    return total_rr / len(queries_relevant) if queries_relevant else 0.0

# 示例
relevant = {2, 4}
retrieved = [1, 3, 2, 5, 4]
print("MRR@5:", mrr_at_k(relevant, retrieved, k=5))   # 1/3 ≈ 0.333

计算 Recall@k：给定真实相关文档 ID 集合和系统返回的 top‑k 文档 ID 列表，计算相关文档被检出的比例。¶

原理

Recall@k 衡量的是所有相关文档中有多少被包含在返回的前 k 个结果中。计算公式：

对于多个查询，通常取平均。

实现

def recall_at_k(relevant_ids: set, retrieved_ids: List[int], k: int) -> float:
    if not relevant_ids:
        return 1.0  # 如果没有相关文档，召回率定义为1还是0？通常设为1.0
    retrieved_set = set(retrieved_ids[:k])
    return len(retrieved_set & relevant_ids) / len(relevant_ids)

def average_recall(queries_relevant: List[set], queries_retrieved: List[List[int]], k: int) -> float:
    total = 0.0
    for rel, ret in zip(queries_relevant, queries_retrieved):
        total += recall_at_k(rel, ret, k)
    return total / len(queries_relevant)

计算 NDCG@k：提供真实相关性分数（如 3,2,1）以及系统返回的文档 ID 及其分数，计算归一化折损累计增益。¶

原理

NDCG（Normalized Discounted Cumulative Gain）考虑了相关性的等级（如 0～3 分）以及文档在结果列表中的位置。计算步骤：

DCG@k：对前 k 个文档，每个文档的相关性分数除以位置的对数折损。

IDCG@k：理想情况下的 DCG，即将所有相关文档按照真实相关性降序排列后计算 DCG。
NDCG@k = DCG@k / IDCG@k。

实现

import math

def ndcg_at_k(relevance_scores: Dict[int, float], retrieved_ids: List[int], k: int) -> float:
    """
    Args:
        relevance_scores: {doc_id: real_score} 真实相关性分数，分数越高越相关。
        retrieved_ids: 系统返回的文档ID列表（按排名顺序）
        k: 截断值
    Returns:
        NDCG@k (float)
    """
    # DCG
    dcg = 0.0
    for i, doc_id in enumerate(retrieved_ids[:k], start=1):
        rel = relevance_scores.get(doc_id, 0.0)
        dcg += rel / math.log2(i + 1)
    # IDCG
    sorted_rels = sorted(relevance_scores.values(), reverse=True)
    idcg = 0.0
    for i in range(1, min(k, len(sorted_rels)) + 1):
        idcg += sorted_rels[i-1] / math.log2(i + 1)
    if idcg == 0.0:
        return 0.0
    return dcg / idcg

# 示例
relevance = {1: 3, 2: 2, 3: 3, 4: 0}   # 文档1和3是最相关的
retrieved = [1, 4, 2, 3]
print("NDCG@4:", ndcg_at_k(relevance, retrieved, k=4))

实现简单的查询扩展：利用 WordNet 或给定的同义词词典，对查询中的词添加同义词，生成多个查询变体以合并检索结果。¶

原理查询扩展通过扩充原始查询的语义范围来提高召回率。这里我们使用 WordNet（通过 nltk）获取同义词，或者使用一个简单的自定义同义词词典。基本流程：对查询分词，为每个词查找同义词，生成多个扩展后的查询（如每个原词加上其同义词的 OR 组合）。为了避免查询漂移，可以只选择最常用的1-2个同义词，或者直接生成一个包含原词和同义词的查询字符串（用 OR 连接），然后用于 BM25 检索（BM25 本身支持 OR 逻辑）。对于稠密检索，可以将多个变体的检索结果合并。

实现

import nltk
from nltk.corpus import wordnet as wn
from typing import List, Set, Optional

nltk.download('wordnet', quiet=True)
nltk.download('omw-1.4', quiet=True)

class QueryExpander:
    def __init__(self, synonym_dict: Optional[dict] = None):
        """
        可选的自定义同义词字典，格式 {word: [syn1, syn2, ...]}
        如果没有提供，则使用WordNet。
        """
        self.custom_dict = synonym_dict or {}

    def get_synonyms(self, word: str, max_syns: int = 2) -> List[str]:
        """获取单词的同义词列表"""
        if word in self.custom_dict:
            return self.custom_dict[word][:max_syns]
        # 使用WordNet
        synonyms = set()
        for syn in wn.synsets(word):
            for lemma in syn.lemmas():
                name = lemma.name().replace('_', ' ').lower()
                if name != word:
                    synonyms.add(name)
        # 返回最常用的前几个
        return list(synonyms)[:max_syns]

    def expand_query(self, query: str, max_syns_per_word: int = 2) -> List[str]:
        """
        返回多个查询变体。这里采用生成一个包含所有同义词的扩展查询字符串（OR形式），
        以及原始查询，都放入列表。也可以只返回一个合并后的查询。
        本函数返回一个列表，每个元素是查询字符串，检索时合并结果。
        """
        words = query.strip().split()
        expanded_words = []
        for w in words:
            # 原词加上其同义词
            syns = self.get_synonyms(w, max_syns_per_word)
            # 构建一个子查询，例如 "word OR syn1 OR syn2"
            # 但是不同的检索器对OR的支持不同，这里我们直接生成多个独立查询变体
            expanded_words.append((w, syns))
        # 生成多个查询变体：原始查询、以及每个替换了同义词的变体
        variants = [query]  # 原始查询
        # 我们可以为每个词生成一个替换了同义词的变体
        for i, (original, syns) in enumerate(expanded_words):
            for syn in syns:
                new_words = words.copy()
                new_words[i] = syn
                variants.append(' '.join(new_words))
        return variants

# 示例
expander = QueryExpander()
variants = expander.expand_query("happy child")
print(variants)
# 输出类似: ['happy child', 'happy kid', 'happy nestling', ...]

结合检索的伪代码

def retrieval_with_expansion(expander, retriever, query, top_k=10):
    variants = expander.expand_query(query)
    all_results = []
    for var in variants:
        results = retriever.search(var, top_k=top_k)  # 调用BM25或稠密检索
        all_results.append(results)
    # 合并多组结果：可以使用RRF融合，或简单合并去重
    merged = merge_results(all_results)   # 自定义合并逻辑
    return merged[:top_k]

说明

查询扩展虽然能提高召回，但可能导致主题漂移，因此通常只添加少量高度相关的同义词。
WordNet 覆盖有限，对于特定领域可以构建自定义同义词表。
在实际系统中，查询扩展常在检索之前进行，生成多个查询，并行检索后融合结果。

实现基于嵌入的查询扩展：获取查询的稠密向量，在预训练的词向量矩阵中检索与其最相似的若干词，扩展原查询。¶

原理

基于嵌入的查询扩展利用预训练词向量（如Word2Vec、GloVe或BERT的静态词表嵌入）在语义空间中寻找与查询词最接近的词汇，将其作为同义词或相关词加入原查询，以扩大检索范围、提高召回率。该方法相比基于词典的同义词扩展，能捕捉更深层的语义关联。

流程：

加载预训练的词向量矩阵（形状 [V, D]）和对应的词表映射。
对查询进行分词，获取每个查询词的向量（若不在词表中则忽略）。
对每个查询词向量，在词向量矩阵中通过余弦相似度检索最相似的 k 个词（排除自身）。
将检索到的相似词加入原查询，生成扩展后的新查询（例如，将原词和相似词用空格拼接）。

实现

import numpy as np
from sklearn.metrics.pairwise import cosine_similarity

class EmbeddingQueryExpander:
    def __init__(self, word_vectors: np.ndarray, word_to_idx: dict, idx_to_word: dict):
        """
        word_vectors: (V, D) 预训练词向量矩阵
        word_to_idx: {word: index}
        idx_to_word: {index: word}
        """
        self.word_vectors = word_vectors
        self.word_to_idx = word_to_idx
        self.idx_to_word = idx_to_word
        # 向量归一化，便于用内积代替余弦相似度
        self.normed_vectors = word_vectors / np.linalg.norm(word_vectors, axis=1, keepdims=True)

    def get_similar_words(self, word: str, top_k: int = 3):
        """返回与给定词最相似的top_k个词（排除自身）"""
        if word not in self.word_to_idx:
            return []
        idx = self.word_to_idx[word]
        query_vec = self.normed_vectors[idx:idx+1]  # (1, D)
        # 计算余弦相似度
        sims = np.dot(query_vec, self.normed_vectors.T).flatten()  # (V,)
        # 排序，取top_k+1（排除自身）
        top_indices = np.argsort(sims)[::-1][1:top_k+1]  # 跳过自己
        similar_words = [self.idx_to_word[i] for i in top_indices]
        return similar_words

    def expand_query(self, query: str, top_k_per_word: int = 2) -> str:
        """扩展查询字符串"""
        words = query.strip().split()
        expanded_tokens = []
        for w in words:
            expanded_tokens.append(w)
            similar = self.get_similar_words(w, top_k_per_word)
            expanded_tokens.extend(similar)
        return ' '.join(expanded_tokens)

示例

# 假设我们有预训练的50维词向量，词表为['car', 'auto', 'vehicle', 'bike', ...]
# 此处仅为示意，实际需加载真实词向量
word_list = ['car', 'auto', 'vehicle', 'bike', 'motorcycle', 'bus']
V = len(word_list)
D = 50
np.random.seed(0)
word_vectors = np.random.randn(V, D)
word_to_idx = {w:i for i,w in enumerate(word_list)}
idx_to_word = {i:w for i,w in enumerate(word_list)}

expander = EmbeddingQueryExpander(word_vectors, word_to_idx, idx_to_word)
new_query = expander.expand_query("fast car", top_k_per_word=2)
print(new_query)  # 可能会输出 "fast car auto vehicle"

说明

实际应用中，词向量应选用大规模预训练模型（如GloVe 6B 300d）或直接使用BERT词表嵌入。
为避免查询漂移，可限制扩展词的数量，或只选用相似度高于某阈值的词。
扩展后的查询可用于BM25或稠密检索，前者直接作为查询字符串，后者可对扩展后的查询重新编码。

实现自适应权重混合检索：根据查询长度或是否包含关键词，自动调整稠密检索与稀疏检索的融合权重，并进行线性融合。¶

原理

混合检索常采用固定权重融合稠密和稀疏分数，但不同查询的最佳权重可能不同。例如，短查询（如“心脏病”）关键词匹配更重要，应提高BM25权重；长查询（如“如何治疗心脏病的最新方法”）语义理解更重要，应提高稠密检索权重。自适应权重策略可根据查询长度、是否包含数字/实体等特征动态调整权重。

实现思路：定义若干规则或训练一个简单的线性模型，根据查询特征预测权重 α（稠密权重），稀疏权重为 1-α。然后对两个检索器返回的文档分数进行Min-Max归一化，再线性加权融合。

实现

from typing import List, Dict, Tuple

def adaptive_weight(query: str, base_alpha: float = 0.5) -> float:
    """根据查询特征自适应调整稠密权重 α"""
    words = query.strip().split()
    query_len = len(words)
    # 规则：长查询提高稠密权重，短查询降低
    if query_len <= 2:
        alpha = max(0.1, base_alpha - 0.2)
    elif query_len >= 6:
        alpha = min(0.9, base_alpha + 0.3)
    else:
        alpha = base_alpha

    # 检查是否包含关键词（数字、专有名词等）
    has_keyword = any(w.isdigit() or w[0].isupper() for w in words)
    if has_keyword:
        alpha -= 0.15  # 提高稀疏权重
    return max(0.0, min(1.0, alpha))

def minmax_normalize(scores: List[float]) -> List[float]:
    if not scores:
        return []
    min_s, max_s = min(scores), max(scores)
    if max_s == min_s:
        return [0.5] * len(scores)
    return [(s - min_s) / (max_s - min_s) for s in scores]

def adaptive_hybrid_retrieve(query: str, dense_scores: Dict[int, float],
                             sparse_scores: Dict[int, float]) -> List[Tuple[int, float]]:
    """
    自适应权重融合两个检索器的分数
    Args:
        dense_scores: {doc_id: cosine_score}  稠密检索返回的文档分数
        sparse_scores: {doc_id: bm25_score}   稀疏检索返回的文档分数
    Returns:
        [(doc_id, fused_score)] 降序排列
    """
    alpha = adaptive_weight(query)
    all_ids = set(dense_scores.keys()) | set(sparse_scores.keys())
    ids = list(all_ids)
    raw_dense = [dense_scores.get(doc, 0.0) for doc in ids]
    raw_sparse = [sparse_scores.get(doc, 0.0) for doc in ids]
    norm_dense = minmax_normalize(raw_dense)
    norm_sparse = minmax_normalize(raw_sparse)
    fused = {}
    for i, doc_id in enumerate(ids):
        fused[doc_id] = alpha * norm_dense[i] + (1 - alpha) * norm_sparse[i]
    return sorted(fused.items(), key=lambda x: x[1], reverse=True)

说明

以上规则仅为示例，实际可基于大量查询日志训练一个分类器或回归模型来预测最佳权重。
融合前进行分数归一化至关重要，否则两种分数的尺度差异会导致权重失效。
该策略可集成到混合检索器中，每次查询动态调用。

利用矩阵乘法加速批量向量检索：输入多个查询向量构成的矩阵和文档向量矩阵，一次性计算所有查询与文档的余弦相似度，返回每个查询的 top‑k。¶

原理当有多个查询需要同时检索时，可以将查询向量堆叠成矩阵 Q (M, D)，文档向量矩阵 D (N, D) 预先归一化。通过一次矩阵乘法 S = Q @ D.T 即可得到所有查询-文档对的余弦相似度矩阵 (M, N)。然后对每一行（每个查询）取 top-k 索引和分数。这种方法充分利用了 GPU 的并行计算能力，相比逐个查询计算极大提升了效率。

实现

import numpy as np

def batch_cosine_search(query_matrix: np.ndarray, doc_matrix: np.ndarray,
                        top_k: int = 10, normalize: bool = True):
    """
    批量余弦相似度检索
    Args:
        query_matrix: (M, D)  M个查询向量
        doc_matrix: (N, D)   N个文档向量
        top_k: 每个查询返回的文档数
        normalize: 是否事先归一化。若为True，则假设输入向量已归一化；
                   否则函数内部会进行归一化。
    Returns:
        top_indices: (M, k)  每个查询的top-k文档索引
        top_scores: (M, k)   对应的余弦相似度
    """
    if normalize:
        # 假设已归一化，直接计算内积
        q = query_matrix
        d = doc_matrix
    else:
        q = query_matrix / np.linalg.norm(query_matrix, axis=1, keepdims=True)
        d = doc_matrix / np.linalg.norm(doc_matrix, axis=1, keepdims=True)
    # 相似度矩阵 (M, N)
    sim_matrix = np.dot(q, d.T)
    # 取每行的top-k（用argpartition加速）
    # 此处使用简单排序
    top_indices = np.argsort(-sim_matrix, axis=1)[:, :top_k]
    top_scores = np.take_along_axis(sim_matrix, top_indices, axis=1)
    return top_indices, top_scores

示例

M, N, D = 5, 1000, 256
queries = np.random.randn(M, D)
docs = np.random.randn(N, D)
indices, scores = batch_cosine_search(queries, docs, top_k=10)
print(indices.shape, scores.shape)

说明

若文档矩阵极大（如百万级），全量矩阵乘法会超出内存，此时应使用近似最近邻索引（如FAISS），但原理与本批量计算相同。
确保向量已归一化，可提前离线对文档矩阵做归一化。

使用 ColBERT 完成检索与重排序流水线：第一阶段用双塔向量召回候选文档，第二阶段用 MaxSim 计算精确相似度并重排序。¶

原理

ColBERT 的典型使用分为两个阶段：

召回阶段：使用双塔模型将查询和文档分别编码为单个向量（例如取 [CLS] token 的输出或对 token 向量做平均池化），然后通过近似最近邻（ANN）检索从海量文档中快速召回 top-K 候选（如 K=1000）。
重排序阶段：对召回的候选文档，加载其预先存储的 token 级向量（所有 token 的嵌入），与查询的 token 级向量进行 MaxSim 计算，得到更精确的相似度分数，最后重新排序输出最终结果。

实现

# 假设已有文档库中所有文档的 token 级向量存储（列表形式，每个文档为 [Ld, D] 张量）
# 以及文档的双塔向量（用于召回）

class ColBERTPipeline:
    def __init__(self, doc_embeddings_dual: np.ndarray, doc_token_vectors: list,
                 query_encoder_dual, query_encoder_token):
        """
        doc_embeddings_dual: (N, D) 文档的双塔向量（已归一化）
        doc_token_vectors: list of (Ld_i, D) tensors，每个文档的token级向量
        query_encoder_dual: 将查询字符串编码为双塔向量的函数
        query_encoder_token: 将查询字符串编码为token级向量的函数
        """
        self.doc_dual = doc_embeddings_dual
        self.doc_tokens = doc_token_vectors
        self.encode_dual = query_encoder_dual
        self.encode_token = query_encoder_token

    def recall(self, query: str, k_candidates: int = 100) -> List[int]:
        """第一阶段：双塔召回，返回候选文档索引"""
        q_vec = self.encode_dual(query)  # (D,)
        q_vec = q_vec / np.linalg.norm(q_vec)
        scores = np.dot(self.doc_dual, q_vec)  # (N,)
        top_indices = np.argsort(-scores)[:k_candidates]
        return list(top_indices)

    def rerank(self, query: str, candidate_indices: List[int]) -> List[Tuple[int, float]]:
        """第二阶段：ColBERT MaxSim 重排序"""
        q_token_vecs = self.encode_token(query)  # (Lq, D)
        scores = {}
        for idx in candidate_indices:
            doc_tokens = self.doc_tokens[idx]    # (Ld, D)
            # 计算 MaxSim（复用之前的 colbert_maxsim 函数）
            score = colbert_maxsim(q_token_vecs, doc_tokens)  # 调用第9题的函数
            scores[idx] = score
        return sorted(scores.items(), key=lambda x: x[1], reverse=True)

    def search(self, query: str, k_final: int = 10, k_candidates: int = 100):
        candidates = self.recall(query, k_candidates)
        ranked = self.rerank(query, candidates)
        return ranked[:k_final]

说明

实际 ColBERT 文档 token 向量存储需要优化（如使用压缩和分页索引），否则占用内存巨大。
双塔向量可通过对 token 向量进行平均池化或取 CLS 得到。
该流水线在保证精度的同时大幅提升检索速度，是 ColBERT 经典部署方式。

实现文档分块检索后的结果合并：若同一文档的多个不同块被检索到，则以文档为单位聚合分数（如取最高块分数），返回文档级排序。¶

原理

在 RAG 系统中，长文档通常被切分为多个块（chunks），每个块独立进行检索。当多个块属于同一文档时，需要在最终的文档级别进行聚合，避免重复文档占据多个排名位置。常用聚合策略：

MaxP：取文档所有块中的最高分数作为文档分数。
MeanP：取平均分。
SumP：求和。这里实现 MaxP 策略，并返回文档级排序（去重且保留最高分）。

实现

from collections import defaultdict
from typing import List, Tuple

def merge_chunk_results(chunk_results: List[Tuple[str, float]], doc_id_func):
    """
    chunk_results: [(chunk_id, score), ...]  chunk_id 可以包含文档ID信息
    doc_id_func: 函数，从 chunk_id 提取所属文档ID
    返回: [(doc_id, max_score)] 降序
    """
    doc_max_scores = defaultdict(float)
    for chunk_id, score in chunk_results:
        doc_id = doc_id_func(chunk_id)
        if score > doc_max_scores[doc_id]:
            doc_max_scores[doc_id] = score
    # 按最高分数降序排列
    return sorted(doc_max_scores.items(), key=lambda x: x[1], reverse=True)

# 示例：chunk_id 格式为 "doc1_chunk0"
def extract_doc_id(chunk_id: str) -> str:
    return chunk_id.split('_')[0]

results = [("doc1_chunk0", 0.8), ("doc1_chunk1", 0.9), ("doc2_chunk0", 0.7), ("doc2_chunk1", 0.6)]
merged = merge_chunk_results(results, extract_doc_id)
print(merged)  # [('doc1', 0.9), ('doc2', 0.7)]

说明

合并后需要重新排序，保留全局 top-k。
在实际应用中，也可以保留每个文档的最高分块的内容，以便生成答案时引用具体片段。

构建学习排序（LTR）的特征向量：对于给定的查询‑文档对，计算其 BM25 分数、嵌入余弦相似度、词重叠比例等特征，组成特征向量。¶

原理

学习排序（Learning to Rank）通过训练一个模型（如LambdaMART）来融合多种信号，预测文档的相关性。第一步是为每个查询-文档对构建特征向量。常用特征包括：

BM25 分数
稠密向量余弦相似度（或双塔模型得分）
词重叠特征（Jaccard 相似度、查询词在文档中的覆盖率等）
文档质量特征（如PageRank、长度）
查询与标题的 BM25 分数等。实现一个特征提取器，输入查询字符串和文档内容，输出一个数值列表。

实现

from sklearn.feature_extraction.text import CountVectorizer
from sklearn.metrics.pairwise import cosine_similarity
import numpy as np

class LTRFeatureExtractor:
    def __init__(self, bm25_retriever, dense_encoder):
        self.bm25 = bm25_retriever   # 具有 score(query, doc_text) 方法
        self.encoder = dense_encoder  # 返回向量的函数

    def extract_features(self, query: str, doc_text: str) -> List[float]:
        features = []
        # 1. BM25 分数
        bm25 = self.bm25.score(query, doc_text)
        features.append(bm25)
        # 2. 稠密余弦相似度
        q_vec = self.encoder(query)
        d_vec = self.encoder(doc_text)
        cos_sim = np.dot(q_vec, d_vec) / (np.linalg.norm(q_vec) * np.linalg.norm(d_vec))
        features.append(cos_sim)
        # 3. 词重叠特征
        query_words = set(query.lower().split())
        doc_words = set(doc_text.lower().split())
        jaccard = len(query_words & doc_words) / len(query_words | doc_words) if query_words | doc_words else 0.0
        features.append(jaccard)
        # 覆盖率：文档中包含的查询词比例
        coverage = len(query_words & doc_words) / len(query_words) if query_words else 0.0
        features.append(coverage)
        # 4. 文档长度（可归一化）
        features.append(len(doc_text.split()))
        return features

说明

实际 LTR 特征可多达数十甚至数百维，包含各种统计和语义特征。
特征应做归一化或标准化，以便模型训练。
该特征向量可输入到 LightGBM、XGBoost 等模型中训练排序模型。

实现一个简单的检索缓存层：以查询字符串的哈希作为键，缓存其 top‑k 检索结果；新查询命中缓存则直接返回，否则执行检索并写入缓存。¶

原理

缓存可以显著减少重复查询的计算开销。实现一个基于内存的 LRU 缓存（或简单字典），以查询字符串的哈希值为键，存储对应的检索结果列表。考虑到查询可能相似但不完全相同，这里采用精确匹配缓存。对于频繁查询的系统，可设置最大容量和过期时间。

实现

import hashlib
from collections import OrderedDict
from typing import List, Tuple, Callable

class RetrievalCache:
    def __init__(self, max_size: int = 1000):
        self.cache = OrderedDict()
        self.max_size = max_size

    def _hash(self, query: str) -> str:
        return hashlib.md5(query.encode('utf-8')).hexdigest()

    def get(self, query: str) -> List[Tuple[int, float]]:
        key = self._hash(query)
        if key in self.cache:
            # 移动到末尾（最近使用）
            self.cache.move_to_end(key)
            return self.cache[key]
        return None

    def put(self, query: str, results: List[Tuple[int, float]]):
        key = self._hash(query)
        if key in self.cache:
            self.cache.move_to_end(key)
        else:
            if len(self.cache) >= self.max_size:
                self.cache.popitem(last=False)  # 移除最久未使用的
        self.cache[key] = results

# 使用示例
cache = RetrievalCache(max_size=500)
def cached_retrieve(retriever, query: str, top_k: int = 10):
    cached = cache.get(query)
    if cached is not None:
        return cached
    # 执行实际检索
    results = retriever.search(query, top_k)   # 假设retriever有search方法
    cache.put(query, results)
    return results

说明

哈希算法可选用更快的方式（如 Python 内置 hash），但 MD5 更稳定，避免哈希碰撞。
实际系统中应设定缓存失效时间（TTL），可存储时间戳并定期清理。
对于语义相似的查询，可考虑使用语义缓存（如基于向量相似度的缓存）。

优化交叉编码器重排序的批量推理：将多个候选文档与同一查询拼接成 batch，利用 attention mask 处理不同长度，一次性得到所有文档的相关性分数。¶

原理

交叉编码器重排序通常需要对每个 (查询, 文档) 对单独进行一次模型推理，当候选文档数量较多（如200个）时，逐对推理速度慢。优化方法：将同一个查询与多个文档分别拼接，形成 batch 输入，一次前向传播得到所有对的分数。由于文档长度不一，需要使用 attention mask 来正确处理填充部分。HuggingFace 的 tokenizer 和模型直接支持批量输入。

实现

from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch

class BatchReranker:
    def __init__(self, model_name: str = "cross-encoder/ms-marco-MiniLM-L-6-v2", device: str = None):
        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.model = AutoModelForSequenceClassification.from_pretrained(model_name).to(self.device)
        self.model.eval()

    def rerank_batch(self, query: str, documents: List[str],
                     max_length: int = 512, batch_size: int = 32) -> List[float]:
        """
        批量重排序：返回与documents顺序对应的分数列表
        """
        # 构造 (query, doc) 对列表
        pairs = [(query, doc) for doc in documents]
        # 分批处理（如果candidate过多，可分多个batch）
        all_scores = []
        for i in range(0, len(pairs), batch_size):
            batch = pairs[i:i+batch_size]
            # 使用 tokenizer 批量编码，自动添加 attention_mask
            encoded = self.tokenizer(
                batch,
                padding=True,           # 填充到本batch最长
                truncation=True,
                max_length=max_length,
                return_tensors="pt"
            )
            inputs = {k: v.to(self.device) for k, v in encoded.items()}
            with torch.no_grad():
                logits = self.model(**inputs).logits  # (batch_size, 1) 或 (batch_size,)
                # 如果是二分类，logits shape (batch_size,)，取第一个token的分数或直接squeeze
                scores = logits.squeeze(-1).cpu().tolist()
                all_scores.extend(scores)
        return all_scores

示例

reranker = BatchReranker()
docs = ["Doc A content...", "Doc B longer content..."] * 10  # 多个候选
scores = reranker.rerank_batch("query text", docs, batch_size=8)
# 然后排序
ranked = sorted(zip(docs, scores), key=lambda x: x[1], reverse=True)

说明

批量推理大幅减少 kernel launch 开销，充分利用 GPU 并行性，速度可达逐对推理的数倍。
max_length 需根据模型和显存调整；可使用动态批处理进一步优化。
对于极大候选集（如数千），可以先用更轻量的方式（如双塔粗筛）缩减候选，再用批量交叉编码器精排。