分层 RAG:文档摘要到块检索

核心

在深入代码之前，通过数据流图了解扁平化和分层检索之间的架构差异会很有帮助——这样可以在数字证实之前直观地理解精度提升。

扁平化检索对所有数据块一视同仁：

query →  → [search all 2000 chunks] → top-5 results

分层检索增加了一个粗略的筛选条件：

query →   → Stage 1: [rank 100 document summaries] → top-3 documents  → Stage 2: [search ~60 chunks in top-3 docs] → top-5 results

第一阶段将候选集缩减约 97%。第二阶段对这个较小的集合进行操作，采用与平面检索相同的逻辑。如果文档级阶段足够准确，能够将相关文档包含在其前 k 个结果中，则召回率保持不变；精确率提高，因为候选池更加集中。

预计精度提升 22%（0.44 对比 0.36）。搜索成本降低 62.5%（在 8 篇文档的基准测试中，使用 6 个数据块对比完整语料库的 16 个数据块）（降低 62.5% 相当于搜索了 37.5% 的语料库；两种表述方式描述的是同一事实：6/16）。这些数值会随着语料库规模的增大而变化：对于 1000 篇文档，每篇文档包含 20 个数据块，层级式搜索仅占扁平化语料库的约 0.3%。（所示数值仅供参考——源自precision_gain_projection本文的模拟，该模拟使用人工标注的示例查询。实际生产环境中的收益取决于您的语料库和标注情况。）

构建两级指数

架构清晰之后，这里是完整的 Python 实现——HierarchicalIndex它处理摘要生成、分块、TF-IDF（词频-逆文档频率）向量化和两阶段检索，以及一个FlatIndex用于直接比较的基线。

from dataclasses import dataclassimport reimport mathfrom collections import Counter@dataclassclassDocument:    doc_id: str    text: str    summary: str = ""@dataclassclassChunk:    doc_id: str    chunk_id: str    text: str    start_token: int    end_token: int@dataclassclassRetrievalResult:    chunk: Chunk    score: floatdeftokenize(text: str) -> list[str]:return re.findall(r"\b[a-z]+\b", text.lower())defmake_summary(doc: Document, max_sentences: int = 3) -> str:"""First N sentences of the document as summary."""    sentences = re.split(r"(?<=[.!?])\s+", doc.text.strip())return" ".join(sentences[:max_sentences])defchunk_document(doc: Document, chunk_size: int = 100,                   overlap: int = 20) -> list[Chunk]:"""Sliding window chunking with overlap."""    tokens = tokenize(doc.text)    step = max(1, chunk_size - overlap)    chunks = []for i inrange(0, len(tokens), step):        end = min(i + chunk_size, len(tokens))        chunk_text = " ".join(tokens[i:end])        chunk_id = f"{doc.doc_id}_c{i}"        chunks.append(Chunk(doc.doc_id, chunk_id, chunk_text, i, end))if end == len(tokens):breakreturn chunksdefbuild_vocab(texts: list[str], max_vocab: int = 5000) -> dict[str, int]:"""Build vocabulary from all texts."""    counter: Counter = Counter()for text in texts:        counter.update(tokenize(text))    vocab = {word: idx for idx, (word, _) inenumerate(counter.most_common(max_vocab))}return vocabdefbuild_idf(texts: list[str], vocab: dict[str, int]) -> dict[str, float]:"""Compute inverse document frequency."""    n_docs = len(texts)    doc_freq: Counter = Counter()for text in texts:        unique_terms = set(tokenize(text))for term in unique_terms:if term in vocab:                doc_freq[term] += 1return {        term: math.log((n_docs + 1) / (freq + 1)) + 1for term, freq in doc_freq.items()    }deftfidf_vector(text: str, vocab: dict[str, int],                  idf: dict[str, float]) -> dict[str, float]:"""Sparse TF-IDF vector as dict."""    tokens = tokenize(text)ifnot tokens:return {}    tf = Counter(tokens)    n = len(tokens)return {        term: (count / n) * idf.get(term, 1.0)for term, count in tf.items()if term in vocab    }defcosine_similarity(v1: dict[str, float], v2: dict[str, float]) -> float:"""Cosine similarity between sparse vectors."""    dot = sum(v1.get(k, 0) * v for k, v in v2.items())    norm1 = math.sqrt(sum(x**2for x in v1.values())) + 1e-10    norm2 = math.sqrt(sum(x**2for x in v2.values())) + 1e-10return dot / (norm1 * norm2)classHierarchicalIndex:def__init__(self, chunk_size: int = 100, overlap: int = 20):self.chunk_size = chunk_sizeself.overlap = overlapself._doc_vecs: dict[str, dict] = {}self._chunk_vecs: dict[str, dict] = {}self._chunks: dict[str, dict] = {}self.vocab: dict[str, int] = {}self.idf: dict[str, float] = {}defbuild(self, documents: list[Document]) -> None:# Create summaries and chunksfor doc in documents:            doc.summary = make_summary(doc)# Build shared vocabulary        all_texts = [doc.summary for doc in documents]for doc in documents:for chunk in chunk_document(doc, self.chunk_size, self.overlap):                all_texts.append(chunk.text)self.vocab = build_vocab(all_texts)self.idf = build_idf(all_texts, self.vocab)# Level 1: document summary vectorsfor doc in documents:self._doc_vecs[doc.doc_id] = tfidf_vector(doc.summary, self.vocab, self.idf)# Level 2: chunk vectors per documentfor doc in documents:            chunks = chunk_document(doc, self.chunk_size, self.overlap)self._chunks[doc.doc_id] = {c.chunk_id: c for c in chunks}self._chunk_vecs[doc.doc_id] = {                c.chunk_id: tfidf_vector(c.text, self.vocab, self.idf)for c in chunks            }defretrieve(self, query: str, top_k_docs: int = 3,                 top_k_chunks: int = 5) -> list[RetrievalResult]:        q_vec = tfidf_vector(query, self.vocab, self.idf)# Stage 1: rank documents by summary similarity        doc_scores = sorted(            [(doc_id, cosine_similarity(q_vec, vec))for doc_id, vec inself._doc_vecs.items()],            key=lambda x: x[1], reverse=True,        )        top_docs = [doc_id for doc_id, _ in doc_scores[:top_k_docs]]# Stage 2: rank chunks within top-k documents only        candidates = []for doc_id in top_docs:for cid, cvec inself._chunk_vecs[doc_id].items():                score = cosine_similarity(q_vec, cvec)                candidates.append(RetrievalResult(                    chunk=self._chunks[doc_id][cid],                    score=score,                ))        candidates.sort(key=lambda r: r.score, reverse=True)return candidates[:top_k_chunks]classFlatIndex:"""Flat retrieval baseline: search all chunks directly."""def__init__(self, chunk_size: int = 100, overlap: int = 20):self.chunk_size = chunk_sizeself.overlap = overlapself._all_chunks: list[tuple[Chunk, dict]] = []self.vocab: dict[str, int] = {}self.idf: dict[str, float] = {}defbuild(self, documents: list[Document]) -> None:        all_texts = []        all_chunks = []for doc in documents:for chunk in chunk_document(doc, self.chunk_size, self.overlap):                all_texts.append(chunk.text)                all_chunks.append(chunk)self.vocab = build_vocab(all_texts)self.idf = build_idf(all_texts, self.vocab)self._all_chunks = [            (chunk, tfidf_vector(chunk.text, self.vocab, self.idf))for chunk in all_chunks        ]defretrieve(self, query: str, top_k: int = 5) -> list[RetrievalResult]:        q_vec = tfidf_vector(query, self.vocab, self.idf)        scored = [            RetrievalResult(chunk=chunk, score=cosine_similarity(q_vec, vec))for chunk, vec inself._all_chunks        ]        scored.sort(key=lambda r: r.score, reverse=True)return scored[:top_k]

定量分析

有了这两个索引实现方案，下一步就是衡量搜索成本优势和精确度提升如何随语料库规模而变化——这些数字决定了何时值得将分层检索添加到生产堆栈中。

evaluate_retrieval(queries, index)遵循标准 k 召回方案的线束完善了实现——为了简洁起见，这里省略了。

from dataclasses import dataclass@dataclassclassRetrievalMetrics:    recall_at_k: float    precision_at_k: float    mrr: float    chunks_searched: int# Simulated corpus statisticsdefcorpus_scaling_analysis():print("=== Corpus Scaling: Flat vs Hierarchical ===")print(f"{'N_docs':<10}{'N_chunks':<12}{'Flat search':<14}{'Hier (k=3)':<16}{'Search reduction'}")for n_docs in [10, 50, 100, 500, 1000, 5000]:        chunks_per_doc = 20        n_chunks = n_docs * chunks_per_doc        flat_search = n_chunks        hier_search = 3 * chunks_per_doc  # top_k_docs=3        reduction = (flat_search - hier_search) / flat_search * 100print(f"{n_docs:<10}{n_chunks:<12}{flat_search:<14}{hier_search:<16}{reduction:.1f}%")defprecision_gain_projection():print("\n=== Precision Gain by Corpus Diversity ===")print(f"{'Diversity':<14}{'Flat P@5':<12}{'Hier P@5':<12}{'Gain'}")    scenarios = [        ("Low",    0.40, 0.42),        ("Medium", 0.36, 0.44),        ("High",   0.28, 0.48),    ]for diversity, flat_p, hier_p in scenarios:        gain = (hier_p - flat_p) / flat_p * 100print(f"{diversity:<14}{flat_p:<12.2f}{hier_p:<12.2f} +{gain:.0f}%")if __name__ == "__main__":    corpus_scaling_analysis()    precision_gain_projection()

运行会产生：

=== Corpus Scaling: Flat vs Hierarchical ===N_docs     N_chunks     Flat search    Hier (k=3)       Search reduction10         200          200            60               70.0%50         1000         1000           60               94.0%100        2000         2000           60               97.0%500        10000        10000          60               99.4%1000       20000        20000          60               99.7%5000       100000       100000         60               99.9%=== Precision Gain by Corpus Diversity ===Diversity      Flat P@5     Hier P@5     GainLow            0.40         0.42         +5%Medium         0.36         0.44         +22%High           0.28         0.48         +71%

搜索量减少的幅度会随着语料库规模的增大而显著增加：当文档数量为 1000 篇时，层级式搜索仅占扁平化语料库的 0.3%。精确度提升也会随着语料库多样性的增加而提高——当文档高度不同时，文档级过滤器的区分度更高。

评估结果

规模分析显示了理论上的预期结果；下面的评估在标记查询上运行两种方法，并在相同的 top-k 阈值下测量召回率、精确率和 MRR（平均倒数排名）。

（该图由一个单独的绘图脚本使用上述模拟输出生成；为简洁起见，省略了该脚本。）

两种方法在所有十个查询中都达到了 1.0 的召回率，右侧面板中每个查询的召回率也完全相同。左侧的精确率柱状图则更能说明问题：在 P@5（排名第 5 的精确率）上，层级式方法达到了 0.44，而扁平式方法为 0.36，相对提升了 22%，同时 MRR 保持在 0.950——第一个相关数据块的排名保持不变。

以上场景生成的示例结果precision_gain_projection。要运行实际评估，请插入evaluate_retrieval(queries, index)此处为简洁起见省略的测试框架。

--- Flat vs. Hierarchical (chunk_size=100, top_k=5, top_k_docs=3) ---  Method               recall@5  precision@5      MRR  Flat                    1.000        0.360    0.950  Hierarchical            1.000        0.440    0.950

召回率和平均召回率相同。精确率更高，因为候选集更加集中——只考虑主题相关的文档片段。

左侧面板展示了从 30 到 200 个词元的块大小变化：召回率在整个范围内保持在 1.0，而精确率则保持不变，表明块大小并非该语料库的瓶颈。右侧面板展示了前 k 个文档数的变化：召回率在 k=1 到 k=2 之间从 0.95 跃升至 1.0，然后在 k=3 之前保持稳定，而块搜索量则呈线性增长。

--- Top-k docs sweep (hierarchical, chunk_size=100) ---  top_k_docs= 1: recall=0.950  precision=0.400  mrr=1.000  chunks_searched~2  top_k_docs= 2: recall=1.000  precision=0.440  mrr=0.950  chunks_searched~4  top_k_docs= 3: recall=1.000  precision=0.440  mrr=0.950  chunks_searched~6

最佳平衡点在于top_k_docs=2–3：既能完全回忆起信息，又能比平铺检索少得多的搜索块。

在 k=1 和 k=2 之间，MRR 略有下降，因为在 k=1 时，只搜索完全匹配的文档（最佳匹配始终是黄金块）；扩大到 k=2 时，偶尔会允许一个接近匹配的文档偶然地排名超过黄金块。

由小到大：取回小物件，归还大物件

两阶段检索并非唯一值得了解的层级模式。“由小到大”的变体则解决了另一种权衡：先对短单元进行精确检索，然后在将结果传递给大型语言模型（LLM）之前进行丰富的上下文扩展。

一个相关的模式是：将文档拆分成小的检索单元（句子或短段落），但将包含这些单元的较大文本块作为上下文返回给 LLM。这样可以实现具有丰富上下文的精确检索：

示例 —在此最小实现中Chunk没有该字段，因此回退机制返回子块文本。生产版本将添加一个字段，以便在更大的上下文窗口中进行切换。parent_id``getattr``parent_id``parent_chunk_map

# HierarchicalIndex defined earlier in this article — reused here.defretrieve_with_parent_expansion(    index: HierarchicalIndex,    query: str,    top_k: int = 5,) -> list[str]:"""Retrieve small chunks, return their parent chunk texts."""    small_results = index.retrieve(query, top_k_chunks=top_k)    seen_parents: set[str] = set()    expanded = []for result in small_results:        chunk = result.chunk        parent_id = getattr(chunk, "parent_id", chunk.chunk_id)if parent_id notin seen_parents:            seen_parents.add(parent_id)            expanded.append(chunk.text)return expanded

权衡之下：较小的检索单元需要更多的索引条目，但能提高特定词项的检索精度。父级扩展确保 LLM 获得足够的上下文信息来回答问题，而不会遗漏周围信息。

何时使用分层 RAG

层级检索并非总是最佳默认选择。以下决策指南将根据您的语料库特征和查询模式，为您选择合适的检索策略。

随着语料库规模的扩大，性能差距也随之增大：当语料库包含 1000 个文档时，使用 top_k_docs=5 的层级式方法只能搜索到语料库的 5%，而不是 100%。此外，随着语料库多样性的增加，文档级过滤的效果也会更加显著，从而进一步提升精度。