数据截至 (上游 commit 5215e9791133)
RSE 相关段落抽取:把碎片拼回长段(查询主算法)
30 秒导读: 普通 RAG 检索返回一堆固定长度的 top-k 碎片,答案却常常横跨连续好几个 chunk。dsRAG 的招牌算法 RSE(Relevant Segment Extraction,相关段落抽取) 换了个思路:先给每个 chunk 打"相关性分",再用一个带约束的优化,动态挑出分数最高的一段连续 chunk——长度不定,能长能短,自动贴合问题的复杂度。项目名里的 "d(ynamic)s(egment)RAG" 精神就落在这一章。
本章讲 dsRAG 的查询侧。摄取侧(把文件切成带上下文头的 chunk)见 01-ingestion-pipeline 与 02-autocontext;组件如何插拔见 04-pluggable-components;更上层的对话与引用见 05-chat-and-citations。
1. 这是什么(先建直觉)
1.1 固定 top-k 的毛病
先看普通 RAG 怎么做检索:把问题嵌入成向量,去向量库里找最相似的前 k 个 chunk(比如 k=5),把这 5 段拼起来喂给 LLM。
问题在于:k 是拍脑袋定死的,而答案的"体量"是不定的。
- 问"公司 2023 年的净利润是多少?"——答案可能就在一个 chunk 里,取 5 个纯属浪费、还塞进噪声。
- 问"详细讲讲这份合同里关于知识产权的全部条款"——答案可能连续横跨 12 个 chunk,取 5 个直接答不全。
更糟的是,固定 top-k 取回的往往是分散的碎片:第 3 段、第 47 段、第 102 段……它们各自相似度高,但彼此不连续,LLM 拿到一堆上下文断裂的片段,很难拼出完整答案。
1.2 RSE 的答案:动态找"连续段"
RSE 的核心主张只有一句:
不要返回 k 个分散的 chunk,而是返回若干条"连续的、最相关的段"(segment),每段的长度由数据自己决定。
一句类比:普通 top-k 像用镊子从书里夹出 5 个词;RSE 像用荧光笔在书上划出几个连续的段落——划多长,取决于这一段有多相关、相关性延续了多久。
一段连续 chunk [start, end) 就叫一个 segment。RSE 要解决的是:在所有 chunk 里,哪几条连续区间的"总相关性"最高,同时不能太长、不能重叠、不能跨到别的文档去。
1.3 用起来什么样
对使用者而言,RSE 是完全隐藏在 query() 后面的:
# 示意:调用方只管发问题,RSE 在内部自动完成
results = kb.query(
search_queries=["知识产权条款有哪些", "IP ownership clauses"], # 可给多个 query
rse_params="balanced", # 预设:balanced / precision / find_all
)
for seg in results:
print(seg["doc_id"], seg["chunk_start"], seg["chunk_end"], seg["score"])
print(seg["content"]) # 已经是拼好的一整段长文本(或 page images)
# 重点看:返回的是「段」,chunk_start~chunk_end 是一段连续区间,不是孤立 chunk
返回的每个 seg 就是一条 segment,content 是这段连续 chunk 拼接起来的整块文本。真实签名见 knowledge_base.py:858 的 KnowledgeBase.query。
2. 顶层全景(查询一次都经过什么)
查询分两大阶段:先把候选 chunk 找出来并排好序(检索/重排),再在这些排序结果上跑 RSE 优化。
怎么读下面这张图:从上到下是数据流,左边是阶段名,右边是产物。
用户: search_queries=[q1, q2, ...] + rse_params
│
┌─────┴──────────────────────────────────────────┐
│ 阶段一:检索 + 重排(每个 query 并行) │ KnowledgeBase._get_all_ranked_results
│ │ → 内部对每个 q 调 _search
│ q ──嵌入──▶ VectorDB.search(top 200) │
│ ──▶ Reranker.rerank_search_results │
│ ──▶ 该 query 的一串 ranked_results │
└─────┬──────────────────────────────────────────┘
│ all_ranked_results = [[...], [...], ...] ← 每个 query 一串
▼
┌─────────────────────────────────────────────────┐
│ 阶段二:RSE(rse.py) │
│ │
│ ① 拼 meta-document ── get_meta_document │ 首尾相接候选文档,记 document_splits
│ ② 给每个 chunk 打分 ── get_relevance_values │ exp(-rank/decay)*sim - penalty
│ ③ 挑最优连续段 ── get_best_segments │ 贪心/约束优化,不重叠、不跨文档
└─────┬──────────────────────────────────────────┘
│ best_segments=[(start,end), ...] + scores
▼
┌─────────────────────────────────────────────────┐
│ 阶段三:取回内容 ── _get_segment_content_... │ text / page_images / dynamic
└─────┬──────────────────────────────────────────┘
▼
list[dict]:{doc_id, chunk_start, chunk_end, content, score, ...}
各步骤的落点:
| 步骤 | 干什么 | 符号(文件) |
|---|---|---|
| 单 query 检索+重排 | 嵌入→向量搜 200→reranker 重排 | _search(knowledge_base.py) |
| 多 query 并行 | 线程池,每 query 一串结果 | _get_all_ranked_results(knowledge_base.py) |
| 拼元文档 | 候选文档首尾相接,记边界 | get_meta_document(rse.py) |
| 打相关性分 | 给元文档每个 chunk 一个分 | get_relevance_values / get_chunk_value(rse.py) |
| 长度归一 | 按 chunk 字符数缩放分数 | adjust_relevance_values_for_chunk_length(rse.py) |
| 挑最优段 | 约束优化选连续区间 | get_best_segments(rse.py) |
| 取回内容 | 三种返回模式 + 段头 | _get_segment_content_from_database(knowledge_base.py) |
后面各节逐个拆开。
3. 查询前半:把候选 chunk 找齐并排序
RSE 的原料,是"每个 chunk 带着一个排名(rank)和一个相似度(similarity)"。这份原料由查询前半段生产。
3.1 单个 query 走一遍:_search
knowledge_base.py:787 的 _search 是最小检索单元,三步:
- 把 query 嵌入成向量(
input_type="query")。 - 向量库搜 top_k 个候选。
- 若非空,交给 reranker 重排,返回重排后的列表。
真实代码(knowledge_base.py:787-797 _search):
query_vector = self._get_embeddings([query], input_type="query")[0]
search_results = self.vector_db.search(query_vector, top_k, metadata_filter)
if len(search_results) == 0:
return []
search_results = self.reranker.rerank_search_results(query, search_results)
return search_results
这里的关键是两段排序:向量搜索先按嵌入相似度粗筛,reranker 再用更强的交叉编码器精排。RSE 后面用到的 rank(名次)和 similarity 都来自重排之后的结果。可插拔的向量库/reranker 见 04-pluggable-components。
3.2 多个 query 并行:_get_all_ranked_results
一次 query() 可以传多个 search query(比如同一意图的中英文两版、或从对话里改写出的几条子问题)。knowledge_base.py:799 的 _get_all_ranked_results 用线程池并行跑,每个 query 各拿一串结果:
with concurrent.futures.ThreadPoolExecutor() as executor:
futures = [executor.submit(self._search, query, 200, metadata_filter)
for query in search_queries]
all_ranked_results = []
for future in futures:
all_ranked_results.append(future.result())
注意两点:
- top_k 在这里被写死为 200(
_search(query, 200, ...))。也就是说不管预设怎么调,每个 query 都先取回 200 个候选 chunk 供 RSE 挑选。 - 返回的
all_ranked_results是**"列表的列表":外层每项对应一个 query,内层是该 query 的 ranked chunk。后面 RSE 的所有函数都吃这个结构,并按 query 分别打分**——这是多 query 能协同扩段的基础。
4. 构造 meta-document:把候选文档首尾相接
RSE 要在"连续 chunk"上做优化,但候选 chunk 可能来自好几个不同文档。为了统一编址,dsRAG 先把涉及的文档首尾相接拼成一条虚拟的长文档——meta-document(元文档),再在这条元文档的全局索引上做优化。
4.1 为什么要拼
设想命中了两个文档:doc A 有 5 个 chunk(04),doc B 有 3 个 chunk(02)。把它 们接起来,就得到一条长度 8 的元文档:
meta-document 全局索引: 0 1 2 3 4 | 5 6 7
对应文档: A0 A1 A2 A3 A4 | B0 B1 B2
↑
document_splits=[5, 8] ← 每个文档「非包含式」的结束点
有了元文档,一段 segment 就能用一对全局索引 [start, end) 表示;而 document_splits(上例 [5, 8])记住每个文档的边界——优化时绝不允许一段 segment 跨过某个 split,否则就会把 A 的尾巴和 B 的开头误拼成一段。
4.2 get_meta_document 在做什么
rse.py:71 的 get_meta_document 产出三样东西:
document_splits:每个文档在元文档里的非包含结束索引(即下一个文档的起点)。document_start_points:每个doc_id的起始全局索引(字典)。unique_document_ids:参与拼接的去重文档 id 列表。
它先从每个 query 的前 top_k_for_document_selection 条结果里收集 doc_id,去重后决定"哪些文档进元文档"(rse.py:73-76):
top_document_ids = []
for ranked_results in all_ranked_results:
top_document_ids.extend([result["metadata"]["doc_id"]
for result in ranked_results[:top_k_for_document_selection]])
unique_document_ids = list(set(top_document_ids))
然后对每个入选文档,扫描全部 ranked 结果,取该文档命中过的最大 chunk_index,用它推算这个文档在元文档里占多长、边界在哪(rse.py:81-88):
for document_id in unique_document_ids:
max_chunk_index = -1
for ranked_results in all_ranked_results:
for result in ranked_results:
if result["metadata"]["doc_id"] == document_id:
max_chunk_index = max(max_chunk_index, result["metadata"]["chunk_index"])
document_start_points[document_id] = document_splits[-1] if document_splits else 0
document_splits.append(int(max_chunk_index + document_splits[-1] + 1
if document_splits else max_chunk_index + 1))
几个要点:
- 文档在元文档里的"长度"取的是命中过的最大 chunk 索引 + 1,不是文档真实全长——RSE 只需覆盖到候 选出现的范围。
document_start_points[doc]在append之前赋值,拿到的正是"上一个文档的结束点",也就是本文档的起点,逻辑自洽。- 元文档总长就是
document_splits[-1](见knowledge_base.py:1036,meta_document_length = document_splits[-1])。
若最终 document_splits 为空(什么都没检索到),query() 直接返回空列表(knowledge_base.py:1031-1033)。
5. 给每个 chunk 打相关性分
meta-document 定好了坐标系,下一步给元文档里每一个位置打一个"相关性分"——分越高,说明这个 chunk 越该被划进某条 segment。这是 RSE 的灵魂公式。
5.1 直觉:命中的 chunk 加分,没命中的扣分
想象把荧光笔划过整条元文档:
- 被某个 query 命中、且排名靠前的 chunk → 正分(值得划)。
- 没被任何 query 命中的 chunk → 负分(不值得划)。
正是这个"没命中就扣分"制造了张力:一条 segment 是若干连续 chunk 的分数求和,中间夹几个没命中的低分 chunk 会拉低总分,但不一定断开——只要两头够相关,算法愿意"跨过"几个平庸 chunk 把它们包进来,从而把碎片重新连成长段。扣多少分,直接决定段落能长到什么程度。
5.2 单个 chunk 的分:get_chunk_value
rse.py:93 的 get_chunk_value 是核心公式:
rank = chunk_info.get('rank', 1000) # 没命中默认排名 1000(极靠后)
absolute_relevance_value = chunk_info.get('absolute_relevance_value', 0.0) # 没命中默认 0
v = np.exp(-rank / decay_rate) * absolute_relevance_value - irrelevant_chunk_penalty
return v
拆解这个式子:
| 项 | 含义 | 效果 |
|---|---|---|
absolute_relevance_value | 该 chunk 的重排相似度(result["similarity"]) | 越相似基础分越高 |
exp(-rank / decay_rate) | 按名次做指数衰减 | 排名越靠后,权重越小;decay_rate 越大衰减越慢(前排更多 chunk 拿到高权重) |
- irrelevant_chunk_penalty | 每个 chunk 统一扣的常数惩罚 | 把"没命中/弱命中"的 chunk 压到负分 |
没命中的 chunk:rank=1000 使 exp(-1000/decay)≈0,再减去 penalty,得到一个负值——正好把它标成"不值得划进段落"。
5.3 penalty 就是"段长旋钮"
irrelevant_chunk_penalty 是最该理解的一个参数:它直接控制段落被切多长。penalty 大 → 平庸 chunk 扣得狠 → 算法不敢包太多 → 段短;penalty 小 → 容忍更多平庸 chunk 被"跨过" → 段长。源码 docstring(rse.py:95-100)给了经验对照:
irrelevant_chunk_penalty | 典型段长(chunk 数) |
|---|---|
| 0.05 | 很长,20~50 |
| 0.1 | 长,10~20 |
| 0.2 | 中,4~10 |
| 0.3 | 短,2~6 |
| 0.4 | 很短,1~3 |
5.4 铺满整条元文档:get_relevance_values
rse.py:109 的 get_relevance_values 把上面的单点公式,对每个 query 分别铺满整条元文档。
它先给每个 query 建一个长度等于元文档的空列表 all_chunk_info,再把该 query 命中的 chunk 按全局索引填进去(rse.py:115-125):
all_chunk_info = [{} for _ in range(meta_document_length)]
for rank, result in enumerate(ranked_results):
document_id = result["metadata"]["doc_id"]
if document_id not in unique_document_ids:
continue # 没进元文档的文档跳过
chunk_index = int(result["metadata"]["chunk_index"])
meta_document_index = int(document_start_points[document_id] + chunk_index) # 文档内索引→全局索引
all_chunk_info[meta_document_index] = {
'rank': rank, # 重排后的名次
'absolute_relevance_value': result["similarity"],
'chunk_length': len(result["metadata"]["chunk_text"]),
}
没被填到的位置保持空 {},get_chunk_value 拿到空 dict 就走默认值(rank=1000、value=0),自然是负分。最后对每个位置调 get_chunk_value,可选再做长度归一(下节),得到该 query 的一条 relevance 向量。所有 query 各一条,汇成 all_relevance_values。
5.5 长度归一:adjust_relevance_values_for_chunk_length
一个细节:chunk 长度不齐时,长 chunk 天然覆盖更多内容,应该更"值钱"。rse.py:139 的 adjust_relevance_values_for_chunk_length 就按字符数缩放分数,reference_length 默认 700(约等于默认切块设置下的平均 chunk 长度):
bounded_chunk_length = max(chunk_length, reference_length) # 只放大「超过参考长度」的 chunk
adjusted_relevance_values.append(relevance_value * (bounded_chunk_length / reference_length))
注意 max(...) 的作用:短于 700 的 chunk 不缩小(比例被夹到 1),只有长于 700 的 chunk 才按比例放大分数。由 rse_params 的 chunk_length_adjustment 开关控制,预设里默认全开。
6. 挑最优段:get_best_segments
有了每个 query 的 relevance 向量,最后一步是选连续区间:在一堆约束下,挑出总分最高的若干条不重叠 segment。这是 rse.py:3 的 get_best_segments。
6.1 它要满足的 约束
一条 segment [start, end) 合法,必须同时满足(见 rse.py:35-51):
| 约束 | 代码含义 |
|---|---|
| 两端都是正分 | relevance_values[start] ≥ 0 且 relevance_values[end-1] ≥ 0 |
| 单段不超长 | 长度 ≤ max_length(内层 range 上界) |
| 不与已选段重叠 | not (start < seg_end and end > seg_start) |
| 不跨文档边界 | 区间内不含任何 document_splits(不越 meta-document 里的文档接缝) |
| 不超总预算 | total_length + (end-start) ≤ overall_max_length |
| 段值达标 | 段总分 ≥ minimum_value,否则该 query 判"没货了" |
段的"值" = 区间内 relevance 之和(rse.py:52,sum(relevance_values[start:end]))。
6.2 贪心 + 多 query 轮转
算法不是全局最优解,而是贪心 + 轮转:外层 while 直到总长度撑满 overall_max_length,或所有 query 都"没货";内层在当前 query 的 relevance 向量里,暴力枚举所有合法 (start, end),取段值最大的那条加入结果。
怎么读下面这张流程:一圈处理一个 query,命中就收一段、换下一个 query;某 query 找不到达标段就把 它拉黑,全部拉黑就收工。
while 总长度 < overall_max_length:
├─ rv_index 轮到下一个 query(转完一圈回到 0)
├─ 所有 query 都进 bad_rv_indices? ── 是 ─▶ break(收工)
├─ 当前 query 已拉黑? ── 是 ─▶ 跳过,换下一个
│
├─ 在当前 query 的 relevance 向量里暴力找最佳合法段
│ (满足 §6.1 全部约束、段值最大)
│
├─ 没找到 或 段值 < minimum_value?
│ └─ 是 ─▶ 把该 query 加入 bad_rv_indices,换下一个
│
└─ 否则:收下这一段 → best_segments/scores 追加,total_length += 段长
真实主体(rse.py:19-67,精简展示):
while total_length < overall_max_length:
if rv_index >= len(all_relevance_values):
rv_index = 0 # 轮转:转完一圈回到第一个 query
if len(bad_rv_indices) >= len(all_relevance_values):
break # 所有 query 都没货了
if rv_index in bad_rv_indices:
rv_index += 1; continue
# ... 在 all_relevance_values[rv_index] 上暴力枚举 (start,end),取段值最大者 ...
if best_segment is None or best_value < minimum_value:
bad_rv_indices.append(rv_index) # 该 query 判定没货
rv_index += 1; continue
best_segments.append(best_segment)
scores.append(best_value)
total_length += best_segment[1] - best_segment[0]
rv_index += 1
轮转是多 query 协同的关键:每轮只给一个 query 收一段,雨露均沾,避免某个 query 一口气吃满 overall_max_length 预算而挤掉别的 query。
6.3 overall_max_length 随 query 数放宽
回到 query():预设里的 overall_max_length 会按 query 数量动态放大(knowledge_base.py:996-998):
overall_max_length += (len(search_queries) - 1) * overall_max_length_extension
多给几个 query,就多给一点总预算(每多一个 query 加 overall_max_length_extension),让更多子问题各自都能分到段落。
6.4 预设:balanced / precision / find_all
rse.py:151 的 RSE_PARAMS_PRESETS 打包了三套调参风格。注意真实预设名是 balanced / precision / find_all(query() 的 docstring 里写的 "precise"/"comprehensive" 是过时文案,以常量为准),query() 默认用 balanced:
| 参数 | balanced | precision | find_all | 作用 |
|---|---|---|---|---|
max_length | 15 | 15 | 40 | 单段最长 chunk 数 |
overall_max_length | 30 | 30 | 200 | 所有段总长预算 |
minimum_value | 0.5 | 0.7 | 0.4 | 段值门槛(越高越挑) |
irrelevant_chunk_penalty | 0.18 | 0.2 | 0.18 | 段长旋钮(见 §5.3) |
overall_max_length_extension | 5 | 5 | 0 | 每多一 query 加的预算 |
decay_rate | 30 | 30 | 200 | 名次衰减速度 |
top_k_for_document_selection | 10 | 10 | 200 | 入选元文档的文档数 |
chunk_length_adjustment | True | True | True | 是否按 chunk 长度归一 |
三套的取向一目了然:
- precision:门槛最高(
minimum_value=0.7)、penalty 最大,只留最相关的短段,追求准。 - find_all:门槛低、预算巨大(总长 200、单段 40)、
top_k_for_document_selection=200几乎不限文档,decay_rate=200让更多靠后 chunk 也保有权重——目标是"尽量把相关内容都捞出来"。 - balanced:居中的默认值。
query() 会用 balanced 作为兜底填补任何缺失参数(knowledge_base.py:972),所以传 dict 时只需覆盖你在意的几个键。
7. 取回内容:_get_segment_content_from_database
get_best_segments 返回的是元文档全局索引下的 (start, end)。query() 先把它们换算回"文档内相对索引"(减去该文档的 doc_start,knowledge_base.py:1071-1086),再逐段取回真实内容。
取内容由 knowledge_base.py:821 的 _get_segment_content_from_database 完成,支持三种 return_mode:
return_mode
├─ "text" ── 段头 + 逐 chunk 文本拼接
├─ "page_images" ── 查该段起止页 → file_system.get_files 取页面图片
└─ "dynamic" ── 段里任一 chunk 是视觉内容? ─是→ page_images
└否→ text
- text 模式:先放一行 segment header(文档标题 + 摘要,由
_get_segment_header生成,见 02-autocontext),再把chunk_start..chunk_end的 chunk 文本依次拼接(knowledge_base.py:843-848)。段头让 LLM 知道"这段出自哪份文档、讲什么",补回被切碎时丢掉的全局上下文。 - page_images 模式:查该段的起止页码,向
file_system要对应的页面图片路径(供 VLM/多模态使用);若这份文档根本没有页面图片,回退到 text 模式(knowledge_base.py:854-855)。 - dynamic 模式:逐 chunk 检查
is_visual,只要有一个是视觉 chunk 就整段走 page_images,否则走 text(knowledge_base.py:828-841)——让"图文混排文档"按需选择返回形态。
最终 query() 给每个 segment 补上 content、页码、score 等字段返回(knowledge_base.py:1089-1112),其中 text 模式下 content 会同时镜像到 text 键作向后兼容。
8. 边界与局限(诚实说)
- 不是全局最优。
get_best_segments是贪心 + 轮转的启发式,不保证选出理论上总分最高的段组合;它换来的是 简单、快、够用。 top_k=200写死。_get_all_ranked_results里每个 query 固定取回 200 个候选(knowledge_base.py:805),不受预设控制;top_k_for_document_selection只影响"哪些文档进元文档",不影响向量搜索深度。- 文档"长度"按命中范围估。 元文档里每个文档的跨度取自"命中过的最大 chunk 索引"(
rse.py:86-88),不是文档真实全长;没被任何 query 命中的尾部 chunk 不进优化空间。 - 段绝不跨文档。
document_splits约束保证一段不会横跨两个文档,这是正确性需要,但也意味着"答案分散在两份文档的接壤处"时会被拆成两段。 - 预设名文档滞后。
query()docstring 里的 "precise"/"comprehensive" 与真实常量precision/find_all不符,传字符串时以RSE_PARAMS_PRESETS的键为准(传错会在knowledge_base.py:968-969抛ValueError)。
9. 代码地图(导航索引)
| 主题 | 文件 | 符号 |
|---|---|---|
| 查询入口 / 编排 | dsrag/knowledge_base.py | KnowledgeBase.query |
| 单 query 检索+重排 | dsrag/knowledge_base.py | KnowledgeBase._search |
| 多 query 并行(top_k=200) | dsrag/knowledge_base.py | KnowledgeBase._get_all_ranked_results |
| 取回段内容(text/page_images/dynamic) | dsrag/knowledge_base.py | KnowledgeBase._get_segment_content_from_database |
| 段头(文档标题+摘要) | dsrag/knowledge_base.py | KnowledgeBase._get_segment_header |
| 拼元文档 + 记边界 | dsrag/rse.py | get_meta_document |
| 单 chunk 相关性公式 | dsrag/rse.py | get_chunk_value |
| 铺满元文档的相关性 | dsrag/rse.py | get_relevance_values |
| 按 chunk 长度归一 | dsrag/rse.py | adjust_relevance_values_for_chunk_length |
| 约束优化选最优段 | dsrag/rse.py | get_best_segments |
| 预设参数 | dsrag/rse.py | RSE_PARAMS_PRESETS |
同组其它章: index(dsRAG 总览) · 01-ingestion-pipeline(摄取管线) · 02-autocontext(上下文头) · 04-pluggable-components(可插拔组件) · 05-chat-and-citations(对话与引用)