Qdrant 2024年夏季代码计划——Python中的ONNX交叉编码器

在Qdrant 2024年夏季代码计划中，实习生Huong (Celine) Hoang成功将ONNX交叉编码器集成到FastEmbed库中，实现了基于相关性分数的搜索结果重排序功能。这一增强显著扩展了Qdrant生态系统的能力，使开发者能够构建更注重上下文的搜索应用，如问答系统。

项目概述 Qdrant以其向量搜索能力闻名，而本次任务更进一步——引入交叉编码器进行重排序。传统上，FastEmbed库负责生成嵌入向量，但交叉编码器并不生成嵌入，而是根据查询与文档的匹配程度输出一组分数。这种重排序在优化搜索结果、将最相关的答案置顶时至关重要。

项目围绕创建新的输入输出方案展开：从文本数据到分数。为此，Hoang设计了一系列类来支持ONNX模型。她处理的关键模型包括Xenova/ms-marco-MiniLM-L-6-v2、Xenova/ms-marco-MiniLM-L-12-v2和BAAI/bge-reranker，这些模型专为重排序任务设计。值得注意的是，FastEmbed是一个极简主义库：不依赖PyTorch或TensorFlow等重型框架，因此非常轻量，占用存储空间小。

技术挑战

1. 构建新的输入输出方案

FastEmbed已支持嵌入生成，但交叉编码器的重排序意味着需要构建全新的类族。这些模型接受查询和一组文档，返回相关性分数列表。为此，Hoang创建了基类如TextCrossEncoderBase和OnnxCrossEncoder，借鉴了现有文本嵌入模型的设计。她确保新类层次结构对用户友好，用户无需了解底层模型的复杂性即可使用。

1. 交叉编码器的分词处理

交叉编码器需要仔细的分词处理，以区分查询和文档。这通过令牌类型ID实现，帮助模型区分两者。Hoang配置分词器处理成对输入——将查询与每个文档连接并分配相应的令牌类型。高效的分词对于模型性能至关重要，她专门针对ONNX模型进行了优化。

1. 模型加载与集成

项目中最有成就感的部分之一是将ONNX模型集成到FastEmbed库中。ONNX模型需要加载到运行时环境中高效管理计算。虽然PyTorch是此类任务的常见框架，但FastEmbed仅支持ONNX模型，从而保持轻量和高效。Hoang进行了大量测试，确保ONNX模型的性能与PyTorch模型相当，使用户能够信任结果。她还增加了批处理支持，允许用户在不损失速度的情况下重排序大量文档。

1. 调试与代码审查

项目中遇到了模型配置、分词器和测试用例等问题。在导师George Panchuk的帮助下，Hoang解决了这些问题，并加深了对最佳实践的理解，特别是代码可读性、可维护性和风格。一个重要的经验是保持代码组织有序和可维护，并确保整个代码库遵循一致清晰的风格。

1. 测试与验证

为确保模型的准确性和性能，Hoang进行了广泛的测试。她将ONNX模型的输出与PyTorch模型的输出进行比较，确保转换正确。关键部分是严格测试以验证输出并识别潜在问题，例如转换错误或实现中的错误。

成果与未来改进 项目结束时，Hoang成功将交叉编码器添加到FastEmbed库中，允许用户基于相关性分数重排序搜索结果。这一增强为依赖上下文排序的应用（如搜索引擎和推荐系统）开辟了新的可能性。该功能随FastEmbed 0.4.0版本发布。

未来改进方向包括：

• 扩展模型支持：添加更多交叉编码器模型，特别是来自sentence transformers库的模型。
• 并行化：优化批处理以处理更大数据集，进一步提升性能。
• 自定义分词：对于非标准分词模型（如BAAI/bge-reranker），添加更具体的分词器配置。

总结与体验 回顾这段经历，Hoang表示这次实习非常宝贵。她不仅在开发技能上有所成长，还学会了如何从头到尾完成复杂项目。Qdrant团队在调试和评审阶段给予了大力支持。她深刻理解了用户体验的重要性——不仅要让模型工作，还要确保它们易于使用并集成到实际应用中。这次经历坚定了她构建有影响力解决方案的热情，并期待未来继续从事此类项目。