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團隊在除錯和評審階段給予了大力支援。她深刻理解了使用者體驗的重要性——不僅要讓模型工作，還要確保它們易於使用並整合到實際應用中。這次經歷堅定了她構建有影響力解決方案的熱情，並期待未來繼續從事此類專案。