RAG 系统性能优化：如何将查询速度提升 10 倍

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13

import chromadb

client = chromadb.PersistentClient(path="./vector_db")

collection = client.get_or_create_collection(
    name="docs",
    metadata={
        "hnsw:space": "cosine",       # 余弦相似度
        "hnsw:M": 32,                 # 图的连通度，越大越准越吃内存
        "hnsw:construction_ef": 200,  # 建索引时的搜索范围
        "hnsw:search_ef": 100,        # 查询时的搜索范围，越大越准越慢
    }
)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33

import jieba
from rank_bm25 import BM25Okapi
import numpy as np

class HybridRetriever:
    def __init__(self, documents):
        self.documents = documents
        # BM25 索引（jieba 中文分词）
        tokenized = [list(jieba.cut(d)) for d in documents]
        self.bm25 = BM25Okapi(tokenized)

    def search(self, query, top_k=5):
        # 路1：BM25 关键词搜索
        tokens = list(jieba.cut(query))
        bm25_scores = self.bm25.get_scores(tokens)
        bm25_rank = np.argsort(bm25_scores)[::-1][:top_k * 3]

        # 路2：向量搜索（接入你自己的向量库）
        # vec_rank = collection.query(query, n_results=top_k*3)
        vec_rank = bm25_rank  # 示例简化

        # RRF 融合两路结果
        return self._rrf_fusion(bm25_rank, vec_rank, top_k)

    def _rrf_fusion(self, rank_a, rank_b, top_k, k=60):
        """RRF 倒数排名融合"""
        scores = {}
        for rank, idx in enumerate(rank_a):
            scores[idx] = scores.get(idx, 0) + 1 / (k + rank + 1)
        for rank, idx in enumerate(rank_b):
            scores[idx] = scores.get(idx, 0) + 1 / (k + rank + 1)
        ranked = sorted(scores.keys(), key=lambda x: scores[x], reverse=True)
        return ranked[:top_k]

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17

from sentence_transformers import CrossEncoder

# 加载交叉编码器
reranker = CrossEncoder('BAAI/bge-reranker-base')

def retrieve_and_rerank(query, collection, top_k=5):
    # 第一阶段：向量粗召回（快，5ms）
    results = collection.query(query_texts=[query], n_results=50)
    candidates = results['documents'][0]

    # 第二阶段：Cross-encoder 精排（准，200ms）
    pairs = [[query, doc] for doc in candidates]
    scores = reranker.predict(pairs)

    # 取分数最高的 top_k
    ranked = sorted(zip(candidates, scores), key=lambda x: x[1], reverse=True)
    return ranked[:top_k]

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11

def hyde_search(query, llm, collection, top_k=5):
    # 步骤1：LLM 生成假设答案
    prompt = f"请简要回答以下问题（2-3句话）：\n{query}"
    hypothetical_answer = llm.generate(prompt)

    # 步骤2：用假设答案做检索（而非原始查询）
    results = collection.query(
        query_texts=[hypothetical_answer],
        n_results=top_k
    )
    return results

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20

import numpy as np

class SemanticCache:
    def __init__(self, embed_fn, threshold=0.95):
        self.embed_fn = embed_fn
        self.threshold = threshold
        self.cache = {}  # {query_emb_tuple: answer}

    def get(self, query):
        query_emb = self.embed_fn([query])[0]
        for cached_key, answer in self.cache.items():
            cached_emb = np.array(cached_key)
            sim = np.dot(query_emb, cached_emb)
            if sim > self.threshold:
                return answer  # 缓存命中，直接返回
        return None  # 未命中

    def set(self, query, answer):
        query_emb = self.embed_fn([query])[0]
        self.cache[tuple(query_emb)] = answer

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17

def parent_child_chunk(documents, chunk_size=200):
    parent_chunks = []
    child_chunks = []

    for doc_id, doc in enumerate(documents):
        # 大块（父）：完整段落，用于最终生成
        parent_chunks.append({"id": f"p_{doc_id}", "text": doc})
        # 小块（子）：按句子拆分，用于精准检索
        sentences = [s.strip() for s in doc.split("。") if len(s.strip()) > 10]
        for i, sent in enumerate(sentences):
            child_chunks.append({
                "id": f"c_{doc_id}_{i}",
                "parent_id": f"p_{doc_id}",
                "text": sent
            })

    return parent_chunks, child_chunks