Hugging Face 推理端点、任务和存储桶如何驱动 Papers with Code 的搜索

来源：https://huggingface.co/blog/pwc-search
三个月前，我们启动了 Papers with Code (https://paperswithcode.co/) 的复兴计划（https://www.reddit.com/r/MachineLearning/comments/1tgmwqr/reviving_paperswithcode_by_hugging_face_p/）（另见公告推文 (https://x.com/NielsRogge/status/2056366395605078252)）。其目标是使开放的 AI 研究易于获取和理解，让人们能轻松找到论文相关的资源，发现 AI 各领域的前沿技术（SOTA），分享有趣的研究，并在彼此的工作基础上进行构建。换言之，其目标是推动引发下一个 Transformer (https://paperswithcode.co/paper/1706.03762) 的研究浪潮。当然，要让 AI 研究易于访问，需要一个强大的搜索引擎，以便人类和智能体能通过网站或 pwc search CLI 命令 (https://github.com/huggingface/pwc-cli) 快速找到相关工作，智能体可通过 Skill (https://github.com/huggingface/pwc-cli/blob/main/standalone_cli/SKILL.md) 使用该命令。

需要注意的是，搜索研究与搜索常规文本有所不同。一个有用的论文搜索引擎应能精确查找标题或 arXiv 标识符，还应理解诸如“用于代码生成的小型语言模型”这样的查询，即使这些词在论文中并未同时出现。它需要识别“原始的 BERT 论文”是一种导航式请求，容忍不完整的标题或拼写错误，并在模型服务冷启动或暂时不可用时仍能快速响应。

Papers with Code 针对查询“DINO”的搜索结果。
Papers with Code 针对查询 DINO 的搜索结果。

对于 Papers with Code (https://paperswithcode.co/)，我们构建了一个混合搜索系统。这也基于我们在 ML6 (http://ml6.eu/) 的先前经验，我们曾在那里为客户开发基于 RAG (https://paperswithcode.co/paper/2005.11401) 的系统。事实证明，混合搜索通常优于基于关键词和向量的搜索系统，因为它结合了两者的优势（另见此博客 (https://techcommunity.microsoft.com/blog/azure-ai-foundry-blog/azure-ai-search-outperforming-vector-search-with-hybrid-retrieval-and-reranking/3929167) 了解更多信息）。关键词搜索找到精确提及内容，而向量搜索找到更模糊、语义相似的术语。请注意，重排器 (https://www.pinecone.io/learn/series/rag/rerankers/)（也称为交叉编码器）可以进一步改善结果，尽管它们也会带来额外的开销和延迟。

显示混合检索优于纯向量检索和关键词搜索的图表。
混合检索优于基于关键词和向量的搜索。图片来自微软《Azure AI Search: Outperforming vector search with hybrid retrieval and reranking》(https://techcommunity.microsoft.com/blog/azure-ai-foundry-blog/azure-ai-search-outperforming-vector-search-with-hybrid-retrieval-and-reranking/3929167) (2023)。

Papers with Code 依赖于 PostgreSQL 数据库，因此其全文搜索功能提供了快速的词法基线。对于稠密嵌入，使用 pgvector (https://github.com/pgvector/pgvector) 添加语义召回，并使用倒数排名融合（RRF）(https://www.elastic.co/docs/reference/elasticsearch/rest-apis/reciprocal-rank-fusion) 算法将两者结合。三个 Hugging Face 服务用于生成稠密嵌入：

• Hugging Face Jobs (https://huggingface.co/docs/hub/jobs) 为我们提供弹性 GPU 计算资源，用于嵌入论文语料库。
• Hugging Face 存储桶 (https://huggingface.co/docs/hub/storage-buckets) 在我们的数据库、实验和 Jobs 之间提供持久化交接。
• Hugging Face 推理端点 (https://huggingface.co/docs/inference-endpoints/index) 为实时查询和增量更新提供低延迟嵌入。

目前，该系统维护着来自 arXiv (https://arxiv.org/) 和 Daily Papers (https://hf.co/papers) 的超过 110,000 篇现有论文的嵌入。本文将解释其背后的架构、设计决策以及我们在将其投入生产过程中所学到的经验。

要点概述

我们特意将搜索分为离线语料库构建和在线搜索服务两部分：
离线语料库构建和在线混合搜索管道架构图。
离线语料库构建和在线混合搜索管道的架构。

计算密集型、面向吞吐量的工作作为 Jobs 运行。持久化制品存储在 Bucket 中。只有小型的查询嵌入步骤位于请求路径上，由一个受保护的推理端点提供支持，以驱动在线搜索。如果该端点冷启动、繁忙或不健康，搜索会立即回退到全文检索。这种分离使系统既强大又快速。

从严格的嵌入契约开始

嵌入管道常以微妙的方式失败：模型版本变更、查询和文档提示混淆、向量截断方式不同，或更新的摘要与其存储的向量不再匹配。我们通过将嵌入格式视为版本化 API 来避免此问题。每篇论文编码为：

归一化标题 + "\n\n" + 归一化摘要

对于每次向量生成，我们记录：

• 模型仓库及确切修订版；
• 输出维度；
• 输入格式版本；
• 输入是查询还是文档；
• 归一化方法；
• 源标题和摘要的内容哈希值。

我们的生产生成使用 Qwen/Qwen3-Embedding-0.6B (https://huggingface.co/Qwen/Qwen3-Embedding-0.6B)，固定到确切修订版，使用 256 维 L2 归一化向量。我们在 MTEB 排行榜 (https://huggingface.co/spaces/mteb/leaderboard) 的帮助下选择了该模型，这是比较嵌入模型的基准。

请注意，像 Qwen3 这样的较新嵌入模型支持两个新功能：

• 可以指定动态嵌入大小，这允许在质量与速度/存储成本之间进行权衡。Qwen 模型称之为“MRL”，即 Matryoshka Representation Learning (https://paperswithcode.co/paper/2205.13147) 的缩写。您可以在此处 (https://huggingface.co/blog/matryoshka) 了解所有信息。我们选择了 256 的嵌入大小以使搜索快速。
• 可以提供指令提示。Qwen 嵌入模型支持 document 提示（我们用它来嵌入论文），实时搜索使用它们的 query 提示（来嵌入用户查询）。

此契约贯穿嵌入从导出、经过 GPU 推理、进入 PostgreSQL，直到在线检索的全过程。

Jobs 将数据库快照转化为向量语料库

全语料库嵌入是典型的批处理工作负载。它需要在相对较短的时间内使用 GPU，受益于高吞吐量，并且不应在运行之间消耗资源。Hugging Face Jobs (https://huggingface.co/docs/huggingface_hub/guides/jobs) 非常适合这种模式：一个 Job 由一个命令、一种硬件配置 (https://huggingface.co/docs/hub/main/en/jobs-pricing#pricing) 和可选的 Docker 镜像定义，可以运行内联声明依赖项的 uv (https://docs.astral.sh/uv/) 脚本。

我们的语料库构建从可重复读 PostgreSQL 快照导出最新版本的每篇论文开始。导出器以流式方式输出行，而不是将目录加载到内存中，写入有界 JSONL 分片，并创建包含行数和 SHA-256 校验和的清单。我们将该不可变运行目录同步到私有存储桶 (https://huggingface.co/docs/hub/storage-buckets)，并使用 hf-mount (https://github.com/huggingface/hf-mount) 直接将 Bucket 挂载到 l4x1 Job（一个具有 24GB 显存的 NVIDIA L4 GPU）中。对于工作者来说，它就是一个文件系统：

hf jobs uv run \
    --flavor l4x1 \
    --timeout 6h \
    --volume hf://buckets/OWNER/pwc-paper-embeddings:/bucket \
    embed_papers_job.py \
    --input /bucket/runs/RUN_ID/input \
    --output /bucket/runs/RUN_ID/output \
    --model Qwen/Qwen3-Embedding-0.6B \
    --revision MODEL_REVISION \
    --dimensions 256 \
    --allow-matryoshka

工作者：

1. 验证输入清单和每个分片的校验和；
2. 加载固定的模型修订版；
3. 按长度对文本排序以减少填充；
4. 批量调用 encode_document（如模型卡片 (https://huggingface.co/Qwen/Qwen3-Embedding-0.6B) 中所述）；
5. 如果 GPU 内存不足，自动减小批次大小；
6. 将 Matryoshka 表示 (https://huggingface.co/blog/matryoshka) 截断到 256 维并进行归一化；
7. 以原子方式写入 float16 Parquet 分片；
8. 记录吞吐量、软件包版本、硬件、峰值显存、行数和输出校验和。

每个完成的分片都有自己的标记，因此重启的 Job 可以跳过已验证的工作。这对于大型语料库很有用：重试应该恢复工作而不是覆盖现有嵌入。在我们 5,000 篇论文的试点中，Qwen Job 在 L4 GPU 上以 1024 维每秒编码约 75 篇论文。相同的处理可以确定性地在 512 和 256 维上实现，因此我们可以比较存储和检索的权衡，而无需支付更多的推理费用。

Bucket 是连接组织

存储桶 (https://huggingface.co/docs/hub/storage-buckets) 是 Hub 上可变的、类似 S3 的对象存储，针对 AI 工作负载进行了优化。它们可以通过 hf://buckets/... 路径访问，并在 Jobs 中无需构建单独的存储集成即可挂载 (https://github.com/huggingface/hf-mount) 进行读写。对我们来说，Bucket 不仅仅是存放向量的地方。它是三个具有不同生命周期的系统之间的边界：

• 生产数据库导出源记录；
• 临时 Jobs 消耗这些记录并产生向量；
• 导入器在接触搜索索引之前验证结果。

我们在不可变的运行前缀下组织制品：

runs/<RUN_ID>/
├── input/
│   ├── manifest.json
│   └── papers-*.jsonl
└── output/
    ├── manifest.json
    ├── embeddings-*.parquet
    └── embeddings-*.complete.json

Bucket 本身是有意设计成可变的，因此不可变性是应用层面的规则：运行 ID 永远不会被覆盖，并且每个制品都由清单和校验和覆盖。这赋予了我们几个有用的属性：

• **可复现性：**我们可以将数据库生成追溯到确切的语料库快照、模型修订版和一组制品。
• **安全重试：**Jobs 可以从同一运行前缀中已完成的分片恢复。
• **低成本实验：**多个模型或维度可以重用一个已验证的输入快照。
• **受控推出：**导入一代不会激活它。我们首先验证覆盖范围并构建其索引。
• **简单回滚：**前一代及其制品保持可用，直到新的一代被证明是稳定的。

只有在导入器重新检查了模式、校验和、维度、归一化、唯一论文 ID 和当前内容哈希之后，我们才将向量加载到 PostgreSQL 中。然后我们为新一代构建一个单独的 HNSW (https://en.wikipedia.org/wiki/Hierarchical_navigable_small_world) 索引，并仅当每一篇符合条件的现有论文都被覆盖时，才以原子方式将其标记为活动状态（HNSW 是实现快速向量搜索的基于图的算法）。

推理端点将语义搜索置于请求路径上

批量嵌入解决了检索的文档侧问题。用户查询仍需要在请求时使用相同的模型契约进行嵌入。我们将固定的模型部署为一个经过身份验证的推理端点 (https://huggingface.co/docs/inference-endpoints/index)，由 Text Embeddings Inference (TEI) (https://github.com/huggingface/text-embeddings-inference) 提供支持。该端点接受查询文本，并使用模型的 query 提示返回一个归一化的 256 维向量。

请注意，此处也可以利用 vLLM (https://github.com/vllm-project/vllm) 或 SGLang (https://github.com/sgl-project/sglang)。

Papers with Code 查询嵌入模型的 Hugging Face 推理端点概述。
用于 Papers with Code 查询嵌入模型的 Hugging Face 推理端点。

然后，API 在活动的 pgvector 代上执行余弦距离搜索：

SELECT paper_id, embedding <=> CAST(:query_vector AS halfvec(256)) AS distance
FROM paper_embeddings
WHERE generation_id = :active_generation
ORDER BY embedding <=> CAST(:query_vector AS halfvec(256))
LIMIT 50;

HNSW 索引使此查找保持快速。在我们的 5,000 篇论文试点中，256 维的 Qwen 索引在 Recall@20 指标上达到了 0.9955（相对于精确搜索），HNSW 查找延迟 p50 为 1.31 毫秒，p95 为 2.21 毫秒。其表和索引占用的存储空间约为 1024 维版本的 27%，同时在该测试中保留了基本相同的 ANN 召回率。

该端点配置为最多一个副本，并且在空闲时可以缩减到零 (https://huggingface.co/docs/inference-endpoints/guides/autoscaling)。这是一个有用的成本杠杆，因为这意味着没有使用时您无需付费。然而，这也意味着冷启动必须是应用程序设计的一部分，而不是被视为异常事件，因为端点启动并服务流量需要一些时间。因此，我们的查询客户端具有刻意严格的控制行为：

• 1 秒的生产超时；
• 非阻塞并发限制；
• 响应维度、有限性和范数验证；
• 基于查询和嵌入生成的短期缓存；
• 在重复失败后的断路器；
• 日志中不记录原始查询文本，仅记录归一化的指纹。

如果端点正在扩展、超时、返回格式错误的向量，或没有可用的并发，我们会立即跳过语义分支。用户仍然会收到词法结果，而不是等待一个不可靠的依赖项。

推理端点运行非常可靠，并且包含一个漂亮的仪表板，以便您快速查看关键分析。

显示请求量、错误、延迟和副本状态的 Hugging Face 推理端点分析仪表板。
显示请求量、错误、延迟和副本状态的推理端点分析。

混合检索比单独任一分支都强大

对于每个查询，词法分支使用加权的 PostgreSQL 全文搜索检索最多 50 个候选。语义分支从 pgvector 检索最多 50 个候选。我们使用加权的倒数排名融合（RRF）结合它们的排名：

score(d) = ∑_{r ∈ {lexical, semantic}} (w_r) / (k + rank_r(d))

RRF 简单且鲁棒，因为它结合的是来自两个不同尺度系统的排名而不是分数。基本上，如果一篇论文在词法分支和语义分支中排名都很高，那么它在混合搜索中排名靠前的机会就更大。我们目前使用相等的分支权重和 k=60（k 是“排名常数”，RRF 算法的一个超参数）。

稠密检索提高了概念性查询的召回率。全文检索在精确术语、标识符和罕见名称方面仍然表现出色。我们还在融合排名之上保留了确定性的身份识别行为：

• 精确标题和 arXiv ID 保持在顶部；
• 方法分类法 (https://paperswithcode.co/methods) 识别导航式搜索，如“原始的 BERT 论文”；
• 不完整的标题和有限的拼写错误使用保守的三元组…