数据截至 (上游 commit dc85934f318c)
04 — 喂进索引的两端
本章讲什么: 文档进 LEANN 之前要过两道工序——切块和算向量。这两道都不难,但踩坑点密集:块太大会被模型截断、块太小丢上下文、批太大爆显存、模板加两遍毁向量。
1. 分块:代码和文本走两条路
1.1 它要解决的小问题
嵌入模型有 token 上限(常见 512),一整个文件塞不进去。所以要切。但切在哪里很影响检索质量:把一个函数从中间劈开,两半都不完整。
1.2 分流
总入口 create_text_chunks(chunking_utils.py:404)按文件扩展名把文档分成两堆:
文档列表
│
▼ detect_code_files:按扩展名判定
├──▶ 代码文件 ──▶ AST 分块(astchunk)──▶ 加行号前缀
│ │失败
│ └──▶ 回退到传统分块
└──▶ 文本文件 ──▶ 传统分块(SentenceSplitter)
判定表 CODE_EXTENSIONS 只认 7 个扩展名 → 4 种语言(chunking_utils.py:131-139):
| 扩展名 | 语言 |
|---|---|
.py | python |
.java | java |
.cs | csharp |
.ts / .tsx / .js / .jsx | typescript |
注意 .js / .jsx 也被当成 typescript 解析。AST 分块只有开了 use_ast_chunking 才走(chunking_utils.py:449)。
1.3 传统分块
用 LlamaIndex 的 SentenceSplitter,分隔符是空格、段落分隔符是 \n\n(chunking_utils.py:346-351)。CLI 里文档默认 256 token / 重叠 128,代码默认 512 / 重叠 50(cli.py:398-418);代码用的那个 splitter 分隔符改成 \n(cli.py:2508-2513)。
有两处防呆:chunk_size <= 0 兜底成 256(chunking_utils.py:314-316),重叠 ≥ 块大小时把重叠砍到一半(:334-344);CLI 侧也独立做了一遍(cli.py:2486-2500)。
1.4 AST 分块
AST(Abstract Syntax Tree,抽象语法树)分块:先把代码解析成语法树,再按函数 / 类这类完整节点的边界切,不会把一个函数劈两半。实现委托给 astchunk 库(仓库的第 6 个 submodule,packages/astchunk-leann)。
三个值得记的细节:
① 输出格式要归一化。 astchunk 可能返回对象、字符串或字典三种形状,_parse_ast_chunk_output 统一成 (text, metadata)(chunking_utils.py:175-193)。
② 单位不一样。 AST 分块的 max_chunk_size 是字符数,传统分块的 chunk_size 是 token 数。换算时按"代码约 1.2 token/字符、约 3 字符/token"的保守估计(calculate_safe_chunk_size,chunking_utils.py:41-70)。
③ 给代码块加行号前缀。 切完之后,如果元数据里有 start_line_no,就给每行加上右对齐的行号和 |(chunking_utils.py:471-479):
42|def search(self, query):
43| return self._index.search(query)
这让检索结果能直接告诉人"在文件的第几行"。
④ 层层回退。 没装 astchunk → 回退传统分块(:209-214);检测不到语言 → 回退(:218-222);解析抛异常 → 回退(:286-289)。
1.5 token 预算
给了 max_tokens_per_chunk 时会做两件事:
- 事前缩小块大小:
safe = 模型上限 * 0.9 - 重叠(calculate_safe_chunk_size,chunking_utils.py:59-64)。 - 事后校验截断:用
tiktoken的cl100k_base编码实测,超 了就截(validate_chunk_token_limits,:73-127)。没装 tiktoken 就退化成按1.2 token/字符的字符截断(:108-118)。
2. 嵌入计算
2.1 总入口与五个 provider
compute_embeddings(embedding_compute.py:422)按 mode 分发:
| mode | 走谁 | 实现 |
|---|---|---|
sentence-transformers(默认) | 本地 HuggingFace 模型 | :502 |
openai | OpenAI 兼容 API | :846 |
mlx | Apple Silicon 的 MLX | :974 |
ollama | 本地 Ollama 服务 | :1051 |
gemini | Google Gemini | :1374 |
CLI 的 --embedding-mode 只开放前四个:build 子命令的 choices 在 cli.py:318-324(:320 那一行),index-* 系列子命令是同一份四选一(cli.py:813-818);嵌入服务自己的 --embedding-mode choices 也是这四个(hnsw_embedding_server.py:515)。也就是说 gemini 只能走 Python API,命令行进不去。
默认模型按平台选:有 NVIDIA GPU 用 BAAI/bge-base-en-v1.5,否则用轻量的 sentence-transformers/all-MiniLM-L6-v2(_default_embedding_model,cli.py:45-64)。
2.2 批大小自适应
设备判定(cuda / mps / cpu)
│
▼
按设备取默认批大小(可用环境变量覆盖)
│
▼ 仅 cuda
按剩余显存再压一次上限
│
▼
encode();若 OOM ──▶ 批大小减半重试 ──▶ 直到 1
- 设备默认值在
_resolve_adaptive_batch_size(:346-352);MPS 上对Qwen/Qwen3-Embedding-0.6B单独降到 32。 - 显存压制在
_cap_cuda_batch_by_vram(:355-380):按"eager attention 峰值内存约 O(seq²)"估算每条序列的开销,只用 20% 空闲显存当预算。 - OOM 重试在
_encode_with_oom_retry(:383-419):捕获 OOM →empty_cache()→ 批大小减半 → 重试,降到 1 还失败才抛。
调用方显式给了 provider_options["batch_size"] 就关掉自适应(:453-456)。
2.3 模型缓存与硬件调优
模型按 (模型名, 设备, 是否 fp16) 组合缓存,避免每批重新加载(:554、:562-566)。加载时按设备做不同调优:CUDA 开 TF32 和 cudnn benchmark;CPU 设线程数并开 mkldnn;MPS 什么都不做(注释解释了原因:set_per_process_memory_fraction 会导致贪婪分配,torch.compile 会撑爆图缓冲区)(:573-592)。
加载参数里 attn_implementation="eager",并优先尝试 local_files_only=True 走本地缓存(:595-621)。
2.4 token 上限探测
get_model_token_limit(:61)是三层回退:
- 运行时缓存,键是
(模型名, base_url)(:58、:80-84)。 - 动态探测:URL 含
11434或ollama走 Ollama 的/api/show;含1234或lmstudio走 LM Studio 的 WebSocket SDK(:87-106)。 - 静态注册表
EMBEDDING_MODEL_LIMITS,支持精确名、去版本号的基名、以及子串模糊匹配三级(:39-53、:110-128)。
都没命中就用默认 2048 并告警(:130-133)。真正截断用 truncate_to_token_limit,前 3 条截断打警告、之后压制(:136-199)。
2.5 prompt 模板:只能对特定模型用
有些嵌入模型(如 EmbeddingGemma)被训练成"文档和查询要加不同前缀"。LEANN 因此支持 build_prompt_template / query_prompt_template,存进 meta 的 embedding_options(api.py:626-627)。
两个必须记住的点:
- 模板只对任务特定的嵌入模型用(
docs/faq.md:14);对普通模型(nomic-embed-text、text-embedding-3-small、bge-base-en-v1.5)加前缀会污染向量,FAQ 原文就是这么写的(docs/faq.md:16)。 - 模板只在
compute_query_embedding里加一次;传给嵌入服务的 provider options 会把三个模板 key 全过滤掉,防止服务端再加一遍(searcher_base.py:79-86、:123-126)。
2.6 平台线程数硬编码
包被 import 时就设环境变量:macOS 上 OMP_NUM_THREADS=1 / MKL_NUM_THREADS=1 / KMP_DUPLICATE_LIB_OK=TRUE;Linux 上同样默认单线程并设 FAISS_NUM_THREADS=1,注释指向 issue #208 的 FAISS/ZMQ 卡死问题(packages/leann-core/src/leann/__init__.py)。这些必须在 import torch 之前生效,所以写在包入口最上面。
3. 本章代码地图
| 主题 | 文件 | 符号 |
|---|---|---|
| 分块总入口 | packages/leann-core/src/leann/chunking_utils.py:404 | create_text_chunks |
| AST 分块 | packages/leann-core/src/leann/chunking_utils.py:196 | create_ast_chunks |
| 传统分块 | packages/leann-core/src/leann/chunking_utils.py:294 | create_traditional_chunks |
| 代码文件判定 | packages/leann-core/src/leann/chunking_utils.py:142 | detect_code_files / CODE_EXTENSIONS |
| 安全块大小换算 | packages/leann-core/src/leann/chunking_utils.py:41 | calculate_safe_chunk_size |
| 块级截断校验 | packages/leann-core/src/leann/chunking_utils.py:73 | validate_chunk_token_limits |
| 嵌入总入口 | packages/leann-core/src/leann/embedding_compute.py:422 | compute_embeddings |
| 本地模型路径 | packages/leann-core/src/leann/embedding_compute.py:502 | compute_embeddings_sentence_transformers |
| 批大小自适应 | packages/leann-core/src/leann/embedding_compute.py:346 | _resolve_adaptive_batch_size |
| 显存压制 | packages/leann-core/src/leann/embedding_compute.py:355 | _cap_cuda_batch_by_vram |
| OOM 减半重试 | packages/leann-core/src/leann/embedding_compute.py:383 | _encode_with_oom_retry |
| token 上限探测 | packages/leann-core/src/leann/embedding_compute.py:61 | get_model_token_limit |
| 截断 | packages/leann-core/src/leann/embedding_compute.py:136 | truncate_to_token_limit |
| provider 配置解析 | packages/leann-core/src/leann/settings.py:224 | encode_provider_options |
| 默认模型选择 | packages/leann-core/src/leann/cli.py:45 | _default_embedding_model |