大家好!我是 Huong (Celine) Hoang,很高兴能分享我今年夏天作为 Qdrant 2024 编程之夏项目一部分的工作经历。在实习期间,我致力于将交叉编码器集成到 FastEmbed 库中,用于重排任务。这一增强扩展了 Qdrant 生态系统的能力,使开发者能够使用 Qdrant 的系列库构建更具上下文感知能力的搜索应用,例如问答系统。
这个项目既具有技术挑战性,也非常有回报,促使我在处理大规模 ONNX(开放神经网络交换)模型集成、Tokenization 等方面提升了技能。让我带您回顾这段旅程、学到的经验以及未来的发展方向。
项目概览
Qdrant 以其向量搜索能力闻名,但我的任务是更进一步——引入交叉编码器进行重排。传统上,FastEmbed 库会生成 embeddings,但交叉编码器不做这个。相反,它们根据查询与文档列表的匹配程度提供一个得分列表。这种重排在您想优化搜索结果并将最相关的答案置顶时至关重要。
该项目围绕创建一个新的输入-输出方案展开:文本数据到得分。为此,我设计了一系列类来支持 ONNX 模型。我使用的一些关键模型包括 Xenova/ms-marco-MiniLM-L-6-v2、Xenova/ms-marco-MiniLM-L-12-v2 和 BAAI/bge-reranker,所有这些都旨在进行重排任务。
需要重点提及的是,FastEmbed 是一个极简库:它没有像 PyTorch 或 TensorFlow 这样的重量级依赖项,因此非常轻量,占用存储空间少得多。
下面是一个图表,展示了该项目的整体工作流程,详细说明了从用户交互到最终输出验证的关键步骤
带重排功能的搜索工作流程

技术挑战
1. 构建新的输入-输出方案
FastEmbed 已经支持 embeddings,但使用交叉编码器进行重排意味着构建一个全新的类族。这些模型接受一个查询和一组文档,然后返回一个相关性得分列表。为此,我创建了基础类,如 TextCrossEncoderBase
和 OnnxCrossEncoder
,并从现有的文本 embedding 模型中获得灵感。
我必须确保的一点是新的类层级结构对用户友好。用户应该能够使用交叉编码器,而无需了解底层模型的复杂性。例如,他们只需要写
同时,在幕后,我们管理所有的模型加载、Tokenization 和得分计算。
from fastembed.rerank.cross_encoder import TextCrossEncoder
encoder = TextCrossEncoder(model_name="Xenova/ms-marco-MiniLM-L-6-v2")
scores = encoder.rerank(query, documents)
2. 处理交叉编码器的 Tokenization
交叉编码器需要仔细处理 Tokenization,因为它们需要区分查询和文档。这通过使用 Token 类型 ID 来完成,这有助于模型区分两者。为了实现这一点,我配置了 tokenizer 来处理成对输入——将查询与每个文档连接起来并相应地分配 token 类型。
高效的 Tokenization 对于确保模型的性能至关重要,我专门针对 ONNX 模型进行了优化。
3. 模型加载与集成
项目中最有成就感的部分之一是将 ONNX 模型集成到 FastEmbed 库中。ONNX 模型需要加载到能够高效管理计算的运行时环境中。
虽然 PyTorch 是处理这类任务的常用框架,但 FastEmbed 仅支持 ONNX 模型,这使其既轻量又高效。我专注于进行大量的测试,以确保 ONNX 模型的性能与其 PyTorch 对应模型相当,从而确保用户可以信任结果。
我还添加了对批处理的支持,允许用户在不牺牲速度的情况下重排大量文档。
4. 调试与代码评审
在项目期间,我遇到了许多挑战,包括模型配置、tokenizer 和测试用例的问题。在我的导师 George Panchuk 的帮助下,我得以解决这些问题,并提升了我对最佳实践的理解,尤其是在代码可读性、可维护性和风格方面。
一个值得注意的经验是保持代码有组织性和可维护性的重要性,并高度关注可读性。这包括正确组织模块并确保整个代码库遵循清晰、一致的风格。
5. 测试与验证
为了确保模型的准确性和性能,我进行了大量的测试。我将 ONNX 模型的输出与其 PyTorch 对应模型的输出进行了比较,确保转换为 ONNX 是正确的。这一过程的关键部分是严格测试,以验证输出并识别潜在问题,例如转换错误或实现中的错误。
例如,验证模型输出的测试结构如下
CANONICAL_SCORE_VALUES
直接取自将原始 PyTorch 模型应用于相同输入的结果
def test_rerank():
is_ci = os.getenv("CI")
for model_desc in TextCrossEncoder.list_supported_models():
if not is_ci and model_desc["size_in_GB"] > 1:
continue
model_name = model_desc["model"]
model = TextCrossEncoder(model_name=model_name)
query = "What is the capital of France?"
documents = ["Paris is the capital of France.", "Berlin is the capital of Germany."]
scores = np.array(model.rerank(query, documents))
canonical_scores = CANONICAL_SCORE_VALUES[model_name]
assert np.allclose(
scores, canonical_scores, atol=1e-3
), f"Model: {model_name}, Scores: {scores}, Expected: {canonical_scores}"
成果与未来改进
在项目结束时,我成功地将交叉编码器添加到了 FastEmbed 库中,允许用户根据相关性得分重排搜索结果。这一增强为依赖上下文排名的应用(如搜索引擎和推荐系统)开启了新的可能性。此功能将从 FastEmbed 0.4.0
版本开始提供。
未来改进的一些方面包括
扩展模型支持:我们可以添加更多交叉编码器模型,特别是来自 sentence transformers 库的模型,为用户提供更多选择。
- 并行化:优化批处理以处理更大的数据集,可以进一步提高性能。
- 自定义 Tokenization:对于具有非标准 Tokenization 的模型,如 BAAI/bge-reranker,可以添加更具体的 tokenizer 配置。
- 总体经验与总结
回顾过去,这次实习是一次极其宝贵的经历。我不仅作为一名开发者成长了,也学会了如何承担复杂的项目并将其从头到尾完成。Qdrant 团队给予了极大的支持,尤其是在调试和评审阶段。我学到了很多关于模型集成、ONNX 以及如何构建用户友好且可扩展的工具的知识。
对我来说,一个重要的收获是理解用户体验的重要性。这不仅仅是让模型正常工作,还要确保它们易于使用并集成到实际应用中。这次经历坚定了我对构建真正有影响力的解决方案的热情,我很高兴将来继续从事这类项目。
感谢您花时间阅读我与 Qdrant 和 FastEmbed 库的这段旅程。我很高兴看到这项工作将如何继续改善用户的搜索体验!
本页面是否有帮助?