构建法律文档高性能向量搜索
我想看看是否可以为一个大型法律数据集构建语义搜索——具体来说，是截至 2023 年澳大利亚法律历史上的所有最高法院判决，分割成 143,485 个可搜索的片段。不是因为有人让我这么做，而是因为规模和领域特定性似乎是一个有趣的挑战。法律文本密集、上下文依赖性强，充满了关键词搜索完全忽略的细微差别。向量搜索能否在大规模下处理这些内容，并保持足够的速度以有用？
我将通过测试不同的嵌入提供商、令人惊讶的性能基准、使用 USearch 和 Isaacus 嵌入 实现此功能的代码，以及最重要的是，在你输入任何你关心的内容之前，为什么要仔细阅读嵌入 API 的条款。
1、选择 API 提供商（或本地运行？）
在深入实现之前，我花了一些时间做每个开发人员都应该做但大多数人没有做的：实际阅读各种嵌入 API 提供商的《服务条款》。我找到的结果……令人清醒。在 Claude 的帮助下，我得到了这张比较表：
更正：我被告知，即使是付费个人用户，Google 也会对你的数据进行训练。对于 Google Workspace 用户（至少在美国），他们不会这样做。
我不是要贬低任何特定的提供商——这些都是合法的企业，有不同的商业模式和限制。但如果你在处理法律文件、医疗记录或任何敏感的信息，这些条款非常重要。一些提供商明确声称有权利用你的数据进行训练，甚至可能在不同条款下与第三方共享或出售数据。其他则需要你手动选择退出，这意味着默认情况下你的数据是可供使用的。在某些情况下，免费层级的用户根本没有保护。
对我来说最突出的是 Isaacus，它只有在你明确提供反馈时才对用户数据进行训练，允许基准测试，并保护免费用户。（完整披露：我与他们合作，所以我可能有点偏见。但即使我没有，这些条款对我来说也很重要。）
2、那么，为什么不直接本地运行呢？
你可能会想：“为什么要处理 API 条款呢？直接在本地运行嵌入模型不就行了。” 这其实是一个完全合理的做法，尤其是对于敏感数据。你可以完全控制，数据不会离开你的基础设施，也不会受制于速率限制或价格变化。
权衡之处在于质量和便利性。一些表现最好的嵌入模型是专有的，只能通过 API 使用。
在这个项目中，我正在测试几种方法。在 API 方面，我比较了 Isaacus、OpenAI、Voyage AI（启用退出选项）和 Google Gemini——都使用它们的最佳模型和最大嵌入维度。我选择了这些是因为它们要么有合理的条款，或者在 Voyage 的情况下，至少允许你选择退出训练（但你必须从一开始就这样做，退出选项不会追溯适用于你已经提交的任何数据）。我还会向你展示一个我微调的小型且快速的本地模型的性能结果：基于 BAAI/bge-small-en 的 sentence-transformers 模型，专门针对澳大利亚最高法院的案例法，来自 Open Australian Legal Corpus……这恰好是我在这里搜索的数据集。
这是否给了我本地模型不公平的优势？绝对如此。它实际上是在我现在查询的数据上训练的。但对于这篇文章，我专注于速度和实现的简便性（建立可靠的本地模型推理基础设施相当复杂）。嵌入模型的选择非常重要，并因用例而异，因此你应该在自己的特定数据上运行自己的评估。我使用的本地模型将句子和段落映射到一个 384 维的密集向量空间，这在大小和速度方面都是合理紧凑的，但不如更大的 API 模型具有语义丰富性。这在向量数据库上的推理型搜索中变得越来越重要。对于更大的开源/本地模型，你很可能在速度和规模上被 API 提供商所超越，至少在这一点上。
3、快速获取嵌入
3.1 异步一切（但小心）
当你处理 143,000 个文本块时，按顺序处理会花费很长时间。所以我采用了异步，但设置了限制以避免过度压榨 API 并遇到速率限制：
# 同时处理最多 5 批
semaphore = asyncio.Semaphore(max_concurrent_batches)
async def process_batch(batch_texts):
async with semaphore: # 速率限制
embeddings = await embed_batch_async(batch_texts)
return embeddings
这个简单的模式给我带来了 3-5 倍的 API 调用速度提升。对于本地模型，异步并没有太大帮助（GPU 计算是瓶颈，而不是网络 I/O），但批处理仍然有效。
3.2 速度数字
在实施这些优化后，以下是不同提供商的实际表现：
批量处理速度（1,000 个法律文档）：
- 本地模型（auslaw-embed，384d）：每秒 924 个文本
- OpenAI（text-embedding-3-large，3072d）：每秒 184 个文本
- Isaacus（kanon-2-embedder，1792d）：每秒 102 个文本
- Google（gemini-embedding-001，3072d）：每秒 19.8 个文本
- Voyage AI（voyage-3-large，2048d）：每秒 14 个文本（哎呀）
单个查询延迟（用户体验）：
- 本地模型：平均 7ms，p95 15ms*
- Google：平均 501.1ms，p95 662.3ms
- OpenAI：平均 1,114ms，p95 1,723ms
- Isaacus：平均 1,532ms，p95 2,097ms
- Voyage AI：平均 1,693ms，p95 7,657ms
注意：这些时间是在澳大利亚工作时间（对我而言相关）。对于海外提供商（OpenAI、Google 和 Voyage），在高峰时段这些时间可以增加约 80%。
*p95 表示“第 95 百分位”——基本上，95% 的请求比这个快。它比平均值更好，因为它显示了你较慢的请求看起来如何，这对用户体验很重要。
这个小型本地模型在速度上彻底击败了 API —— 批次速度快 5 倍，单个查询快 70 倍。但这里有个问题，它使用的是 384 维嵌入，而 API 是 1792–3072 维。这个维度差异对搜索质量很重要，这就是为什么我在这篇文章中专注于 API 提供商。
一些意外发现：
- Voyage 的限流很痛苦：每秒 14 个文本的付费 API 很糟糕
- Google 不受益于异步：似乎还是顺序处理
- OpenAI 非常一致：尽管是基于网络的，但方差最低
- Isaacus 处于很好的中间位置：在维度上具有良好的速度
4、使用 Isaacus 嵌入法律文本
经过测试所有提供商后，我将专注于 Isaacus 的深入研究。kanon-2-embedder 嵌入也具有这样一个特性，即前几个维度携带了大部分信息，使我们能够做一些有趣的事情。
下面是使用 Isaacus 生成嵌入并保存以备以后使用的简单方法。你需要先安装所需的包，并且需要 Python ≥3.8。
pip install isaacus numpy
import asyncio
import numpy as np
from isaacus import AsyncClient
import os
async def generate_embeddings():
"""使用 Isaacus API 为你的法律语料库生成嵌入。"""
# 初始化异步客户端（API 密钥来自环境变量）
client = AsyncClient(api_key=os.getenv("ISAACUS_API_KEY"))
# 你的语料库和查询
corpus_texts = ["The High Court of Australia...", ...] # 143,485 文档
queries = ["What is the highest court?", ...] # 你的查询
# 生成语料库嵌入，使用任务感知编码
corpus_response = await client.embed(
model="kanon-2-embedder",
inputs=corpus_texts,
task="retrieval/document" # 告诉模型这些是文档
)
corpus_embeddings = np.array(corpus_response.embeddings, dtype=np.float32)
# 生成样本查询嵌入，使用任务感知编码
query_response = await client.embed(
model="kanon-2-embedder",
inputs=queries,
task="retrieval/query" # 告诉模型这些是查询
)
query_embeddings = np.array(query_response.embeddings, dtype=np.float32)
# 如果不存在，则创建嵌入目录
os.makedirs("embeddings", exist_ok=True)
# 保存到磁盘以备以后使用
np.save("embeddings/corpus_embeddings.npy", corpus_embeddings)
np.save("embeddings/query_embeddings.npy", query_embeddings)
await client.close()
# 运行它
asyncio.run(generate_embeddings())
就是这样！嵌入现在已保存并准备好进行优化。
5、256 维度的魔法
Isaacus 嵌入是 1792 维的，但它们经过特殊训练，其中前几个维度携带了最多的相关信息（类似于主成分）。这意味着我们可以截断到仅 256 维度，同时仍保持出乎意料的好的搜索质量：
# 加载我们保存的完整嵌入
corpus_embeddings = np.load("embeddings/corpus_embeddings.npy")
# 仅使用前 256 个维度
corpus_256d = corpus_embeddings[:, :256].astype(np.float32)
这导致...
- 8.6 倍更快 的搜索（459 q/s vs 53 q/s）
- 7 倍更少内存（140 MB vs 1,028 MB）
- 61% recall@10†（vs 100% for full dimensions — explained below）
- 57% recall@50
关于这些数字的重要背景： 这些召回百分比是相对于合成基线而言的，而不是绝对检索质量。以下是我是如何做的：
- 我使用相同的文档作为语料库和查询（143K 文档查询自己）
- “完美”的 100% 基线是完整的 1792 维精确搜索——每个文档都将其自身作为顶级结果
- 61% 和 57% 的数字展示了在使用 256 维度进行优化时丢失了多少信息
这是一个优化过程中信息损失的基准测试，不是真实世界的检索质量。对于实际的法律搜索质量，你需要人工标注的测试集（如 MLEB 基准测试）。
为什么这仍然重要： 它告诉你优化保留了大约 60% 的排序顺序。缺失的 40% 并不是错误的结果，只是排名不同。对于 RAG，你无论如何都会拉取 50-100 个块，这通常是可以接受的。
†Recall@10 表示“我们实际找到的真正前 10 名结果的百分比？” Recall@50 更低，因为此测试使用自我匹配的相同文档——随着 k 增大，会有更多“完美”匹配要检索，所以错过任何一个都会降低召回率分数。
6、USearch：每秒2,880 查询（只需 CPU！）
现在是最有趣的部分。大多数向量搜索教程会告诉你使用 FAISS 或 Pinecone 并构建近似索引。但我想要一种不同的方式：在我需要的时候有出色的结果，当我不需要的时候有闪电般的速度，并且最关键的是 不需要 GPU。
进入 USearch。这是一个不太为人知的库，通过 SIMD 优化（基本上是现代 CPU 指令，一次处理多个数据点）实现了精确和近似搜索。其杀手级功能？一切都运行在 CPU 上，这意味着：
- 部署成本更低：不需要昂贵的 GPU 实例
- 更容易扩展：只需添加更多的 CPU 核心
- 更低的复杂性：没有 CUDA 驱动程序，没有 GPU 内存管理
- 仍然非常快：得益于 SIMD 向量化
对于全天候运行的法律搜索系统，避免 GPU 可以将基础设施成本降低 70-80%，同时仍提供亚毫秒级响应时间。这是一场革命。
你可以通过以下方式安装它：
pip install usearch
级别 1：只使用更多核心
# 使用多线程进行批量处理
matches = search(corpus, queries, 100, MetricKind.Cos, exact=True, threads=8)
# 结果：374 q/s（比单线程快 7 倍）
级别 2：构建 HNSW 索引
对于查询多于更新的工作负载，构建索引是有回报的：
from usearch.index import Index
index = Index(
ndim=1792,
metric=MetricKind.Cos,
connectivity=32, # 更高 = 更好质量，更多内存
expansion_add=200, # 构建质量
expansion_search=100 # 搜索质量
)
# 一次性构建（214 秒）
for i, embedding in enumerate(corpus_embeddings):
index.add(i, embedding)
# 多次搜索（飞快）
matches = index.search(query, 100)
# 结果：993 q/s 且 98.6% recall@10
HNSW（Hierarchical Navigable Small World graphs）本质上是一种巧妙的数据结构，让你比暴力搜索更快地找到近似最近邻，同时仍然非常准确。
级别 3：完整栈
结合 256d 减少、HNSW 索引和半精度存储：
# 准备 256d 嵌入，使用半精度
corpus_256d = corpus_embeddings[:, :256].astype(np.float16)
index = Index(
ndim=256,
metric=MetricKind.Cos,
dtype="f16", # 半精度节省 2 倍内存
connectivity=32,
expansion_add=200,
expansion_search=100
)
# 构建（59 秒用于 143K 文档）
for i, emb in enumerate(corpus_256d):
index.add(i, emb)
最终数字：
- 2,880 查询/秒（比基线快 54 倍！）
- 每次查询 0.35 毫秒（亚毫秒！）
- 总内存 70 MB（内存减少 14.7 倍）
- 61% recall@10（仍然足够用于 RAG）
完整的性能阶梯：
提醒：这些是相对于同一语料库上的 1792 维精确搜索而言的。这些衡量的是优化权衡，而不是绝对的检索质量。对于法律特定的检索基准，请参阅 MLEB。
基线系统（53 q/s）：
- 1 个用户：19ms（尚可）
- 100 个并发用户：1.9 秒（不是很好）
- 1,000 个用户：19 秒（真的，真的不好）
优化系统（2,880 q/s）：
- 1 个用户：0.35ms
- 100 个并发用户：35ms
- 1,000 个并发用户：347ms
- 10,000 个并发用户：3.5s
优化后的系统可以在单台机器上处理大量流量。
7、选择你的配置
正确的选择取决于你的需求：
使用“准确性”模式（基线 + 多线程），如果：
- 法律合规要求完美的召回率
- 你进行仔细的研究，而不是 RAG
- 语料库足够小（<100K 文档）
使用“平衡”模式（HNSW，完整维度），如果：
- 你需要接近完美的结果（>95% 召回率）
- 正在构建生产级法律搜索工具
- 能够承担 3-4 分钟的索引构建时间
- 这是我推荐用于法律应用的模式
使用“速度”模式（全栈），如果：
- 构建面向消费者的应用程序
- 内存非常有限（<200MB 用于索引）
- 使用 RAG 与重新排序（60% 召回率足以作为第一遍）
- 需要处理数千个并发用户
8、生产就绪的代码
这里是你可以适应的完整实现：
import numpy as np
from usearch.index import Index, search, MetricKind
from pathlib import Path
from typing import Optional, Union, List
import time
class OptimizedLegalSearch:
def __init__(
self,
corpus_embeddings: np.ndarray,
optimization_level: str = "balanced",
save_path: Optional[str] = None
):
"""
初始化搜索系统。
Args:
corpus_embeddings: 语料库嵌入 (N × D 数组)
optimization_level: 一个 "accuracy", "balanced" 或 "speed"
save_path: 可选路径以保存/加载索引
"""
self.corpus_embeddings = corpus_embeddings
self.optimization_level = optimization_level
self.save_path = save_path
self.index = None
# 根据优化级别配置
if optimization_level == "speed":
# 最大速度: 256d + HNSW + f16
print("配置为最大速度 (256d + HNSW + f16)...")
self.use_dimensions = 256
self.use_index = True
self.dtype = "f16"
self.connectivity = 32
self.expansion_add = 200
self.expansion_search = 100
elif optimization_level == "balanced":
# 平衡: 全维 + HNSW
print("配置为平衡速度/质量 (HNSW M=32)...")
self.use_dimensions = None # 使用所有维度
self.use_index = True
self.dtype = "f32"
self.connectivity = 32
self.expansion_add = 200
self.expansion_search = 100
elif optimization_level == "accuracy":
# 最大准确性: 全维，精确搜索
print("配置为最大准确性 (精确搜索)...")
self.use_dimensions = None
self.use_index = False
self.dtype = "f32"
else:
raise ValueError(f"未知的优化级别: {optimization_level}")
# 准备语料库
self._prepare_corpus()
# 如果需要，构建或加载索引
if self.use_index:
if save_path and Path(save_path).exists():
self.load_index(save_path)
else:
self._build_index()
if save_path:
self.save_index(save_path)
def _prepare_corpus(self):
"""根据优化设置准备语料库嵌入。"""
if self.use_dimensions:
# 截断到指定维度
self.corpus_processed = self.corpus_embeddings[:, :self.use_dimensions]
if self.dtype == "f16":
self.corpus_processed = self.corpus_processed.astype(np.float16)
else:
self.corpus_processed = self.corpus_processed.astype(np.float32)
print(f"语料库缩减到 {self.use_dimensions} 维度")
else:
self.corpus_processed = self.corpus_embeddings.astype(np.float32)
print(f"使用全部 {self.corpus_embeddings.shape[1]} 维度")
# 确保连续内存以便 SIMD 优化
self.corpus_processed = np.ascontiguousarray(self.corpus_processed)
def _build_index(self):
"""构建 HNSW 索引以实现快速近似搜索。"""
ndim = self.use_dimensions or self.corpus_embeddings.shape[1]
print(f"构建 HNSW 索引 (M={self.connectivity}, ef={self.expansion_search})...")
start_time = time.time()
self.index = Index(
ndim=ndim,
metric=MetricKind.Cos,
dtype=self.dtype,
connectivity=self.connectivity,
expansion_add=self.expansion_add,
expansion_search=self.expansion_search
)
# 添加向量到索引并跟踪进度
n_docs = len(self.corpus_processed)
for i, embedding in enumerate(self.corpus_processed):
self.index.add(i, embedding)
if (i + 1) % 10000 == 0:
print(f" 已索引 {i+1}/{n_docs} 个文档...")
build_time = time.time() - start_time
print(f"索引构建耗时 {build_time:.1f} 秒")
# 报告内存使用情况
memory_mb = self.index.size * (2 if self.dtype == "f16" else 4) / (1024 * 1024)
print(f"索引内存使用: {memory_mb:.1f} MB")
def search(
self,
query_embeddings: Union[np.ndarray, List[np.ndarray]],
k: int = 100,
return_scores: bool = False
) -> Union[np.ndarray, List[np.ndarray]]:
"""
搜索最相似的文档。
Args:
query_embeddings: 查询向量(s) - 可以是单个或批量
k: 要返回的结果数量
return_scores: 是否返回相似度分数
Returns:
最相似文档的索引（和可选的分数）
"""
# 确保 numpy 数组
if isinstance(query_embeddings, list):
query_embeddings = np.vstack(query_embeddings)
# 如果需要，截断查询维度
if self.use_dimensions:
if len(query_embeddings.shape) == 1:
query_processed = query_embeddings[:self.use_dimensions]
else:
query_processed = query_embeddings[:, :self.use_dimensions]
else:
query_processed = query_embeddings
# 确保浮点 32 用于查询（更好的精度）
query_processed = query_processed.astype(np.float32)
if self.use_index:
# HNSW 近似搜索
if len(query_processed.shape) == 1:
# 单个查询
matches = self.index.search(query_processed, k)
if return_scores:
return matches.keys, matches.distances
return matches.keys
else:
# 批量查询
results_keys = []
results_scores = []
for q in query_processed:
matches = self.index.search(q, k)
results_keys.append(matches.keys)
if return_scores:
results_scores.append(matches.distances)
if return_scores:
return results_keys, results_scores
return results_keys
else:
# 精确搜索，带多线程
matches = search(
self.corpus_processed,
query_processed,
k,
MetricKind.Cos,
exact=True,
threads=8 # 使用 8 个线程进行并行处理
)
if return_scores:
return matches.keys, matches.distances
return matches.keys
def save_index(self, path: str):
"""将 HNSW 索引保存到磁盘以实现快速加载。"""
if self.index:
print(f"将索引保存到 {path}...")
self.index.save(path)
# 也保存元数据
import json
metadata = {
"optimization_level": self.optimization_level,
"use_dimensions": self.use_dimensions,
"dtype": self.dtype,
"connectivity": self.connectivity,
"expansion_add": self.expansion_add,
"expansion_search": self.expansion_search,
"corpus_size": len(self.corpus_processed)
}
with open(f"{path}.meta.json", "w") as f:
json.dump(metadata, f, indent=2)
print("索引保存成功")
def load_index(self, path: str):
"""从磁盘加载预构建的 HNSW 索引。"""
print(f"从 {path} 加载索引...")
self.index = Index.restore(path)
# 如果可用，加载元数据
import json
meta_path = f"{path}.meta.json"
if Path(meta_path).exists():
with open(meta_path, "r") as f:
metadata = json.load(f)
print(f"加载 {metadata['optimization_level']} 索引，包含 {metadata['corpus_size']} 个文档")
# 示例用法
if __name__ == "__main__":
# 加载你的嵌入
corpus_embeddings = np.load("embeddings/corpus_embeddings.npy")
# 如果存储方式类似，加载你的查询
query_embeddings = np.load("embeddings/query_embeddings.npy")
# 创建所需优化的搜索系统
searcher = OptimizedLegalSearch(
corpus_embeddings,
optimization_level="balanced", # 或 "speed" 或 "accuracy"
save_path="indices/legal_search.index"
)
# 搜索相似文档
results = searcher.search(query_embeddings[0], k=100)
print(f"\n找到了 {len(results)} 个相似文档")
一个最小的查询工作流程可以简单如下：
query = "什么是先例原则？"
query_emb = client.embed(model="kanon-2-embedder", inputs=[query], task="retrieval/query")
results = searcher.search(np.array(query_emb.embeddings[0]), k=5)
print(results)
9、结束语
对于法律搜索，我可能会坚持“平衡”配置——993 q/s 且 98.6% recall@10 足够快，同时保持接近完美的准确性。但对于一般的 RAG 应用，你无论如何都会拉取 50 多个块并重新排序，那么完全优化的“速度”模式在 2,880 q/s 时非常吸引人。
一个 61% 召回率的系统在 0.35ms 内响应通常胜过一个 100% 召回率的系统需要 19ms，特别是当你的检索只是多步骤流水线的第一阶段时。
现在你可以构建自己的高性能法律搜索引擎！
原文链接：How I Built Lightning-Fast Vector Search for Legal Documents
汇智网翻译整理，转载请标明出处