QdrantのSummer of Code 2024プログラムで、インターンのHuong (Celine) Hoangは、FastEmbedライブラリにONNXクロスエンコーダを統合し、関連性スコアに基づく検索結果の再ランク付け機能を実現しました。この強化により、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のような重い依存関係を持たないため、非常に軽量でストレージ容量も小さく済みます。
技術的課題
- 新しい入出力スキームの構築
FastEmbedはすでに埋め込みをサポートしていましたが、クロスエンコーダによる再ランク付けには、完全に新しいクラスファミリーの構築が必要でした。これらのモデルはクエリとドキュメントのセットを受け取り、関連性スコアのリストを返します。そのため、Hoangは既存のテキスト埋め込みモデルから着想を得て、TextCrossEncoderBaseやOnnxCrossEncoderなどのベースクラスを作成しました。ユーザーが基盤となるモデルの複雑さを知らなくてもクロスエンコーダを操作できるよう、クラス階層はユーザーフレンドリーに設計されました。
- クロスエンコーダのトークン化処理
クロスエンコーダでは、クエリとドキュメントを区別するために慎重なトークン化が必要です。これは、モデルが両者を区別するのに役立つトークンタイプIDを使用して行われます。Hoangは、クエリと各ドキュメントを連結し、それに応じてトークンタイプを割り当てるペア入力の処理を行うようにトークナイザを設定しました。効率的なトークン化はモデルのパフォーマンスに不可欠であり、特にONNXモデル向けに最適化されました。
- モデルのロードと統合
プロジェクトで最もやりがいのある部分の一つは、ONNXモデルをFastEmbedライブラリに統合することでした。ONNXモデルは、計算を効率的に管理するランタイム環境にロードする必要があります。PyTorchはこれらのタスクで一般的なフレームワークですが、FastEmbedはONNXモデルのみをサポートすることで軽量かつ効率的を保っています。Hoangは、ONNXモデルがPyTorchモデルと同等のパフォーマンスを発揮することを確認するために広範なテストを実施し、ユーザーが結果を信頼できるようにしました。また、バッチ処理のサポートも追加し、ユーザーが速度を損なうことなく大量のドキュメントを再ランク付けできるようにしました。
- デバッグとコードレビュー
プロジェクト中、モデル設定、トークナイザ、テストケースに関する問題に直面しました。メンターのGeorge Panchukの助けを借りて、これらの問題を解決し、コードの可読性、保守性、スタイルに関するベストプラクティスの理解を深めました。特に重要な教訓は、コードを整理整頓し、保守可能に保ち、明確で一貫性のあるスタイルを維持することの重要性でした。
- テストと検証
モデルの精度とパフォーマンスを確保するため、Hoangは広範なテストを実施しました。ONNXモデルの出力をPyTorchモデルと比較し、変換が正しいことを確認しました。このプロセスの重要な部分は、出力を検証し、変換エラーや実装のバグなどの潜在的な問題を特定するための厳格なテストでした。
成果と今後の改善 プロジェクトの終了時点で、HoangはFastEmbedライブラリにクロスエンコーダを追加することに成功し、ユーザーが関連性スコアに基づいて検索結果を再ランク付けできるようになりました。この機能強化により、検索エンジンやレコメンデーションシステムなど、コンテキストによるランク付けに依存するアプリケーションに新たな可能性が開かれました。この機能はFastEmbed 0.4.0から利用可能です。
今後の改善点としては、以下が挙げられます：
- モデルサポートの拡大：特にsentence transformersライブラリからのクロスエンコーダモデルを追加し、ユーザーにより多くの選択肢を提供する。
- 並列化：バッチ処理を最適化してより大規模なデータセットを処理し、パフォーマンスをさらに向上させる。
- カスタムトークン化：BAAI/bge-rerankerのような非標準トークン化を必要とするモデルに対して、より特化したトークナイザ設定を追加する。
総括と体験 振り返ってみると、このインターンシップは非常に貴重な経験でした。開発者としてだけでなく、複雑なプロジェクトを最初から最後まで遂行できる能力も成長しました。Qdrantチームは、特にデバッグとレビューの段階で多大なサポートを提供してくれました。モデル統合、ONNX、そしてユーザーフレンドリーでスケーラブルなツールの構築について多くのことを学びました。
重要な教訓は、ユーザーエクスペリエンスを理解することの重要性です。単にモデルを動かすだけでなく、それらを実際のアプリケーションに簡単に統合して使用できるようにすることが求められます。この経験は、真にインパクトのあるソリューションを構築することへの情熱を確固たるものにし、今後もこのようなプロジェクトに取り組み続けることを楽しみにしています。