向量相似度检索服务

源文件：chapter3/dense-embedding/README.md

向量相似度检索服务

一个用于教学的 HTTP 向量相似度检索服务，基于 BGE-M3 嵌入模型，可在 ANNOY 与 HNSW 两种索引后端之间灵活切换。

命令行工具：稠密检索与 ANN 对比（cli.py，实验 3-4）

除了上面的 HTTP 服务，本项目还提供一个开箱即用、可离线复现的命令行工具 cli.py，
把书中实验 3-4 的两个观察点直接跑成可量化的数字，无需先启动服务：

1. 稠密嵌入检索的语义能力——在带标注的小型语料上计算 recall@k / precision@k / MRR；
2. ANN 索引后端对比（实验 3-4 的重点）——复用服务端 indexing.py 里的 ANNOY / HNSW
   实现，测量二者相对精确暴力检索的召回率、建索引耗时与查询延迟。

用法

# 1) 单条稠密查询（默认查询 "a cat playing"，需要嵌入模型）
python cli.py -q "model distillation" -k 3

# 2) 检索质量评测：recall@k / precision@k / MRR
python cli.py --eval

# 2') 离线复现：用已缓存的小模型（无需下载 2.3GB 的 BGE-M3）
python cli.py --embedding-model sentence-transformers/all-MiniLM-L6-v2 --eval

# 3) ANN 后端对比（合成向量，完全离线、无需任何模型）
python cli.py --compare-ann -k 10
python cli.py --compare-ann --backend hnsw --hnsw-ef-search 200 -k 10   # 调 ef_search 看召回随之上升

# 自定义语料 / 标注 / 输出
python cli.py --corpus my.json --labels my_labels.json --eval -o result.json
​

python cli.py --help 提供完整的中文参数说明（--corpus / --query / --embedding-model / --top-k / --output，以及 ANN 对比的各项索引超参）。

常用参数

实测结果（真实运行，非杜撰）

稠密检索质量（内置 12 篇语料，all-MiniLM-L6-v2，离线）：

宏平均  recall@5=1.000  precision@5=0.320  MRR=1.000
​

其中查询 a cat playing 的相关文档只用 kitten / feline 表达、不含字面 "cat"，
稠密检索仍把它们排到第 1、2 名——这正是稠密相对稀疏 BM25（实验 3-5 会漏召回）的语义优势。

ANN 后端对比（3000 条 128 维随机单位向量，100 条查询，top-10）：HNSW 的召回率随
ef_search 单调上升，体现"精度 / 速度"取舍：

环境提示：本工具对每个后端会先做"自查询自身向量"的健康检查。在部分 macOS/arm64 环境下，
annoy==1.17.3 的预编译轮子存在缺陷（连查询库中已有向量都只返回它自己），此时工具会打印
[警告] ...疑似当前环境下损坏 并把该后端的数字标记为不可信。HNSW 不受影响。若要复现完整的
ANNOY vs HNSW 对比，请在 annoy 正常工作的环境（如 Linux x86_64）中运行。

特性

• BGE-M3 模型：业界领先的多语言嵌入模型，支持：

  • 用于语义检索的稠密嵌入
  • 多语言支持（100+ 种语言）
  • 长上下文（最长 8192 个 token）

• 双索引后端：

  • ANNOY（Approximate Nearest Neighbors Oh Yeah）：快速、内存高效的基于树的索引
  • HNSW（Hierarchical Navigable Small World）：高精度的基于图的索引

• 教学型日志：详尽的调试日志，展示：

  • 嵌入生成过程
  • 索引操作（插入 / 删除 / 检索）
  • 性能指标
  • 向量统计信息

• RESTful API：简洁的 HTTP 接口，用于：

  • 文档索引（插入 / 更新）
  • 文档删除
  • 相似度检索
  • 统计与监控

• 内存存储：为保持简洁采用纯内存运行（不做持久化）

架构

┌──────────────────┐
│   HTTP Client    │
└────────┬─────────┘
         │
         ▼
┌──────────────────┐
│  FastAPI Server  │
└────────┬─────────┘
         │
         ▼
    ┌────┴────┐
    │         │
    ▼         ▼
┌──────────┐ ┌──────────────┐
│ Document │ │  Embedding   │
│  Store   │ │   Service    │
└──────────┘ │  (BGE-M3)    │
             └──────┬─────────┘
                    │
                    ▼
          ┌─────────┴──────────┐
          │                    │
          ▼                    ▼
    ┌──────────┐        ┌──────────┐
    │  ANNOY   │        │   HNSW   │
    │  Index   │        │  Index   │
    └──────────┘        └──────────┘
​

安装

前置条件

• Python 3.8 或更高版本
• macOS（针对 M1/M2 芯片优化）或 Linux
• 至少 4GB 内存（推荐 8GB）
• 兼容 CUDA 的 GPU（可选，用于加快嵌入生成）

安装步骤

1. 安装依赖：

cd projects/week3/dense-embedding
pip install -r requirements.txt
​

2. 下载 BGE-M3 模型（首次运行时会自动下载）：
   • 模型体积：约 2.3GB
   • 会缓存到 HuggingFace 缓存目录中

使用方法

启动服务

使用 HNSW 索引（默认）

python main.py
​

使用 ANNOY 索引

python main.py --index-type annoy
​

自定义配置

python main.py --index-type hnsw --host 0.0.0.0 --port 4242 --debug --show-embeddings
​

可用选项

• --index-type：选择索引后端（annoy 或 hnsw，默认：hnsw）
• --host：服务监听地址（默认：0.0.0.0）
• --port：服务端口（默认：4240）
• --debug：开启调试模式，输出详细日志
• --show-embeddings：在日志中显示嵌入向量（教学用途）

API 文档

服务启动后，可访问：

• 交互式文档：http://localhost:4240/docs
• OpenAPI schema：http://localhost:4240/openapi.json

API 端点

1. 索引文档

POST /index
​

索引一篇新文档或更新已有文档。

请求：

{
  "text": "Machine learning is a subset of artificial intelligence.",
  "doc_id": "doc_001",  // 可选，未提供时自动生成
  "metadata": {         // 可选元数据
    "category": "AI",
    "author": "John Doe"
  }
}
​

响应：

{
  "success": true,
  "doc_id": "doc_001",
  "message": "Document indexed successfully using hnsw",
  "index_size": 1
}
​

2. 检索文档

POST /search
​

检索相似文档。

请求：

{
  "query": "What is deep learning?",
  "top_k": 5,
  "return_documents": true
}
​

响应：

{
  "success": true,
  "query": "What is deep learning?",
  "results": [
    {
      "doc_id": "doc_001",
      "score": 0.8543,
      "text": "Machine learning is a subset...",
      "metadata": {"category": "AI"},
      "rank": 1
    }
  ],
  "total_results": 5,
  "search_time_ms": 12.5
}
​

3. 删除文档

DELETE /index
​

从索引中删除一篇文档。

请求：

{
  "doc_id": "doc_001"
}
​

响应：

{
  "success": true,
  "message": "Document doc_001 deleted successfully",
  "index_size": 0
}
​

4. 获取统计信息

GET /stats
​

获取服务统计信息。

响应：

{
  "index_type": "hnsw",
  "index_size": 100,
  "document_count": 100,
  "embedding_dimension": 1024,
  "model_name": "BAAI/bge-m3"
}
​

5. 列出文档

GET /documents?limit=10
​

列出存储中的文档。

测试

运行演示客户端

演示客户端用示例文档展示全部功能：

python test_client.py
​

运行性能测试

用合成数据测试索引和检索性能：

python test_client.py --performance
​

使用 curl 手动测试

索引一篇文档：

curl -X POST http://localhost:4240/index \
  -H "Content-Type: application/json" \
  -d '{"text": "This is a test document about machine learning."}'
​

检索相似文档：

curl -X POST http://localhost:4240/search \
  -H "Content-Type: application/json" \
  -d '{"query": "artificial intelligence", "top_k": 5}'
​

索引对比

ANNOY（Approximate Nearest Neighbors Oh Yeah）

优点：

• 索引速度极快
• 内存占用低
• 适合静态数据集
• 支持多种距离度量

缺点：

• 删除需重建索引
• 构建后无法增量更新
• 速度与精度之间存在取舍（由 n_trees 控制）

最适用于：

• 大规模相似度检索
• 读多写少的负载
• 内存受限的环境

HNSW（Hierarchical Navigable Small World）

优点：

• 召回精度高
• 支持增量更新
• 检索速度快且精度好
• 支持软删除

缺点：

• 内存占用更高
• 索引构建比 ANNOY 慢
• 参数调优更复杂

最适用于：

• 动态数据集
• 高精度要求
• 读写均衡的负载

配置

环境变量

可通过以 VEC_ 为前缀的环境变量来配置服务：

export VEC_INDEX_TYPE=hnsw
export VEC_MODEL_NAME=BAAI/bge-m3
export VEC_USE_FP16=true
export VEC_MAX_SEQ_LENGTH=512
export VEC_MAX_DOCUMENTS=100000
export VEC_LOG_LEVEL=DEBUG

# ANNOY 专用
export VEC_ANNOY_N_TREES=50
export VEC_ANNOY_METRIC=angular

# HNSW 专用
export VEC_HNSW_EF_CONSTRUCTION=200
export VEC_HNSW_M=32
export VEC_HNSW_EF_SEARCH=100
export VEC_HNSW_SPACE=cosine
​

教学特性

本服务为教学目的内置了详尽的日志：

1. 嵌入生成日志：展示文本转换为向量的过程
2. 索引操作日志：关于索引更新的详细信息
3. 检索过程日志：逐步展示检索执行过程
4. 性能指标：所有操作的耗时信息
5. 向量统计：嵌入的最小 / 最大 / 均值（启用时）

开启完整的教学日志：

python main.py --debug --show-embeddings
​

性能考量

内存占用

• BGE-M3 模型：约 2.3GB
• 每篇文档开销：约 4KB（1024 维 float32 嵌入）
• ANNOY 索引：约 (4 * dimension * n_items * n_trees / 2) 字节
• HNSW 索引：约 (4 * dimension * n_items * M * 2) 字节

优化建议

1. 针对 ANNOY：

   • 增大 n_trees 可提升精度（但构建更慢）
   • 归一化向量时使用 angular 度量
   • 批量插入后再构建索引

2. 针对 HNSW：

   • 增大 M 可提升召回（但更耗内存）
   • 增大 ef_construction 可提升索引质量（但构建更慢）
   • 调整 ef_search 在速度与精度间权衡

3. 通用：

   • 使用 FP16 加快推理（精度略降）
   • 尽量批量插入文档
   • 根据文档长度合理限制 max_seq_length

故障排查

常见问题

1. 内存不足：

   • 减小批大小
   • 使用 FP16 模式
   • 降低 max_seq_length
   • 用 ANNOY 代替 HNSW

2. 索引过慢：

   • 减小 HNSW 的 ef_construction
   • 减小 ANNOY 的 n_trees
   • 如有 GPU 则启用

3. 检索质量差：

   • 增大 ANNOY 的 n_trees
   • 增大 HNSW 的 M 和 ef_search
   • 检查文档是否过短 / 过长

参考资料

• BGE-M3 论文
• BGE-M3 模型
• ANNOY 库
• HNSWlib
• FastAPI 文档

许可证

本项目为学习用途的教学项目。