Sentence Transformers semantic_search 怎样控制语义检索的分块大小

Sentence Transformers 的 semantic_search 分块大小，直接在调用时设置 query_chunk_size 和 corpus_chunk_size。前者控制一次处理多少条查询向量，后者控制一次扫描多少条语料向量；内存紧张时先调小它们，速度允许时再逐步增大。top_k 只控制每条查询最多返回多少个匹配项，不负责分块。

官方地址：https://sbert.net/docs/package_reference/util/retrieval.html

先把“分块”定位到检索循环

遇到语料库变大后内存突然上升，先不要修改模型或盲目降低 top_k。官方函数签名把输入分为二维的查询向量和语料向量，再按块计算相似度。query_chunk_size 是同时处理的查询数，corpus_chunk_size 是一次扫描的语料数。它们影响计算窗口，不会改变向量本身的维度。

把查询向量、语料向量、query_chunk_size、corpus_chunk_size、相似度函数和 top_k 放在同一张关系图里，便于判断问题属于哪一层：资源峰值通常看两个 chunk，返回数量看 top_k，相似度含义则看 score_function。

用显式参数覆盖默认分块

写成显式配置比依赖默认值更适合服务端。下面的示例只展示调用关系，不声称已经在本机执行；参数名和返回结构按官方文档组织。

from sentence_transformers.util import semantic_search

# query_embeddings 和 corpus_embeddings 都应是二维向量张量
hits = semantic_search(
    query_embeddings,
    corpus_embeddings,
    query_chunk_size=64,       # 一次处理64条查询，先控制查询侧峰值
    corpus_chunk_size=100000,  # 一次扫描10万条语料，避免默认窗口过大
    top_k=8,                   # 每条查询最多返回8个结果，不是分块参数
)

# 每个查询对应一个结果列表，列表内按分数从高到低排列
for query_id, matches in enumerate(hits):
    print(f"查询{query_id}返回{len(matches)}条")
    for item in matches:
        # item 主要包含 corpus_id 与 score，按需映射回原始文本
        print(item["corpus_id"], item["score"])

这里的 64 和 100000 不是通用最佳值，而是一个更保守的起点。查询数量少、语料较小的离线任务可以适当增大；共享 GPU 或并发请求较多的服务，通常先减小语料块，再观察吞吐是否仍能接受。

不要把编码 batch_size 当成检索分块

常见误判是只调整 model.encode(..., batch_size=...)。这个参数影响文本转向量时一次送入模型的文本数量；semantic_search 接收的已经是向量，随后才由两个 chunk 参数控制查询与语料的扫描窗口。两阶段都可能产生峰值，因此应分别记录。

# 编码阶段：batch_size 控制文本进入模型的批量
query_embeddings = model.encode(
    queries,
    batch_size=32,  # 只影响文本到向量的阶段
    normalize_embeddings=True,
)

# 检索阶段：再单独控制查询块和语料块
hits = semantic_search(
    query_embeddings,
    corpus_embeddings,
    query_chunk_size=64,
    corpus_chunk_size=100000,
    top_k=8,
)

# 两个阶段都要保留异常信息，方便区分输入为空与资源不足
if not hits:
    raise ValueError("没有查询结果，请检查查询向量和语料向量")

按现象调整参数并复查结果

可以采用单变量调参：先固定模型、输入、top_k 和相似度函数，只把 corpus_chunk_size 从较小值逐步提高；若查询批很多，再调整 query_chunk_size。记录峰值内存和总耗时，而不是只看单次最快结果。减小 chunk 通常会降低单次资源峰值，但会增加循环次数，最终速度取决于向量维度、设备、并发和数据规模。

大语料场景的边界

官方文档说明该函数适合约百万条以内的语料检索。调小 chunk 只能缓解一次计算窗口的资源压力，不能把线性扫描变成索引检索；当语料继续增长或延迟要求固定时，应评估 FAISS 等专门检索方案，并保持相同的向量、距离函数和 top_k 语义再做对比。

query_chunk_size 越小，结果会变差吗？在输入、相似度函数和精度不变时，它主要改变分批计算方式，不应因分块本身改变排序；但仍应在目标设备上复查边界输入。

把 corpus_chunk_size 改成 0 可以关闭分块吗？不要这样猜。该参数表示每次扫描的语料数量，应使用正整数并以当前版本官方签名为准；需要更大窗口就明确增大数值。

排查时记住三条线：模型编码看 batch_size，检索资源看两个 chunk，结果数量看 top_k。把三者分开记录，才能在内存、耗时和召回结果之间做出可解释的调整。