OpenVINO と Postgres を使用した高速かつ効率的なセマンティック検索システムの構築

Pixabay の real-napster による写真

最近のプロジェクトの 1 つでは、高パフォーマンスで拡張でき、レポート検索に対してリアルタイムの応答を提供できるセマンティック検索システムを構築する必要がありました。これを実現するために、AWS RDS 上で PostgreSQL と pgvector を AWS Lambda と組み合わせて使用​​しました。課題は、ユーザーが厳密なキーワードに依存するのではなく、自然言語クエリを使用して検索できるようにしながら、応答が 1 ～ 2 秒未満、あるいはそれ以下であることを保証し、CPU リソースのみを利用できるようにすることでした。

この投稿では、検索から再ランキングまで、この検索システムを構築するために行った手順と、OpenVINO とトークン化のためのインテリジェントなバッチ処理を使用した最適化について説明します。

セマンティック検索の概要: 取得と再ランキング

現代の最先端の検索システムは通常、検索と再ランキングという 2 つの主要なステップで構成されます。

1) 取得: 最初のステップでは、ユーザーのクエリに基づいて関連ドキュメントのサブセットを取得します。これは、OpenAI の大小のエンベディング、Cohere の Embed モデル、Mixbread の mxbai エンベディングなど、事前トレーニングされたエンベディング モデルを使用して実行できます。検索では、クエリとの類似性を測定することでドキュメントのプールを絞り込むことに重点を置いています。

これは、この目的で私のお気に入りのライブラリの 1 つである Huggingface の文変換ライブラリを検索に使用した簡略化された例です。

from sentence_transformers import SentenceTransformer
import numpy as np

# Load a pre-trained sentence transformer model
model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")

# Sample query and documents (vectorize the query and the documents)
query = "How do I fix a broken landing gear?"
documents = ["Report 1 on landing gear failure", "Report 2 on engine problems"]

# Get embeddings for query and documents
query_embedding = model.encode(query)
document_embeddings = model.encode(documents)

# Calculate cosine similarity between query and documents
similarities = np.dot(document_embeddings, query_embedding)

# Retrieve top-k most relevant documents
top_k = np.argsort(similarities)[-5:]
print("Top 5 documents:", [documents[i] for i in top_k])

2) 再ランキング: 最も関連性の高い文書が取得されたら、クロスエンコーダー モデルを使用してこれらの文書のランキングをさらに改善します。このステップでは、より深い文脈の理解に重点を置き、クエリに関連して各ドキュメントをより正確に再評価します。
再ランキングは、各ドキュメントの関連性をより正確にスコアリングすることで、さらに洗練された層を追加するため、有益です。

これは、軽量クロスエンコーダーであるcross-encoder/ms-marco-TinyBERT-L-2-v2を使用した再ランキングのコード例です:

from sentence_transformers import CrossEncoder

# Load the cross-encoder model
cross_encoder = CrossEncoder("cross-encoder/ms-marco-TinyBERT-L-2-v2")

# Use the cross-encoder to rerank top-k retrieved documents
query_document_pairs = [(query, doc) for doc in documents]
scores = cross_encoder.predict(query_document_pairs)

# Rank documents based on the new scores
top_k_reranked = np.argsort(scores)[-5:]
print("Top 5 reranked documents:", [documents[i] for i in top_k_reranked])

ボトルネックの特定: トークン化と予測のコスト

開発中、文変換のデフォルト設定で 1,000 件のレポートを処理すると、トークン化と予測の段階にかなりの時間がかかることがわかりました。特にリアルタイムの応答を目指していたため、これによりパフォーマンスのボトルネックが発生しました。

以下では、SnakeViz を使用してコードのプロファイリングを行い、パフォーマンスを視覚化しました。

ご覧のとおり、トークン化と予測の手順が不釣り合いに遅く、検索結果の提供に大幅な遅れが生じます。全体的には平均して 4 ～ 5 秒かかりました。これは、トークン化と予測のステップの間にブロック操作があるためです。データベース呼び出し、フィルタリングなどの他の操作も合計すると、合計で簡単に 8 ～ 9 秒かかります。

OpenVINO によるパフォーマンスの最適化

私が直面した質問は次のとおりです: 高速化は可能ですか? 答えは「はい」です。CPU 推論用に最適化されたバックエンドである OpenVINO を活用することで実現できます。 OpenVINO は、AWS Lambda で使用する Intel ハードウェアでのディープ ラーニング モデル推論の高速化に役立ちます。

OpenVINO 最適化のコード例
OpenVINO を検索システムに統合して推論を高速化する方法は次のとおりです:

from sentence_transformers import SentenceTransformer
import numpy as np

# Load a pre-trained sentence transformer model
model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")

# Sample query and documents (vectorize the query and the documents)
query = "How do I fix a broken landing gear?"
documents = ["Report 1 on landing gear failure", "Report 2 on engine problems"]

# Get embeddings for query and documents
query_embedding = model.encode(query)
document_embeddings = model.encode(documents)

# Calculate cosine similarity between query and documents
similarities = np.dot(document_embeddings, query_embedding)

# Retrieve top-k most relevant documents
top_k = np.argsort(similarities)[-5:]
print("Top 5 documents:", [documents[i] for i in top_k])

このアプローチにより、2 ～ 3 倍のスピードアップが得られ、元の 4 ～ 5 秒が 1 ～ 2 秒に短縮されます。完全に動作するコードは Github にあります。

速度の微調整: バッチ サイズとトークン化

パフォーマンスを向上させるもう 1 つの重要な要素は、トークン化 プロセスを最適化し、バッチ サイズ と トークンの長さ を調整することでした。バッチ サイズ (batch_size=16) を増やし、トークンの長さ (max_length=512) を減らすことで、トークン化を並列化し、反復操作のオーバーヘッドを減らすことができます。私たちの実験では、batch_size が 16 から 64 の間でうまく機能し、それより大きいとパフォーマンスが低下することがわかりました。同様に、max_length を 128 に決定しました。これは、レポートの平均長が比較的短い場合に実行可能です。これらの変更により、全体で 8 倍の高速化を達成し、CPU 上でも再ランキング時間を 1 秒未満に短縮しました。

実際には、これは、データの速度と精度の適切なバランスを見つけるために、さまざまなバッチ サイズとトークン長を試してみることを意味しました。これにより、応答時間が大幅に改善され、1,000 件のレポートでも検索システムを拡張できるようになりました。

結論

OpenVINO を使用し、トークン化とバッチ処理を最適化することで、CPU のみのセットアップでリアルタイム要件を満たす高性能のセマンティック検索システムを構築することができました。実際、全体で 8 倍の高速化が実現しました。センテンストランスフォーマーを使用した検索とクロスエンコーダーモデルを使用した再ランキングの組み合わせにより、強力でユーザーフレンドリーな検索エクスペリエンスが作成されます。

応答時間と計算リソースに制約がある同様のシステムを構築している場合は、OpenVINO とインテリジェントなバッチ処理を検討して、パフォーマンスを向上させることを強くお勧めします。

この記事を楽しんでいただければ幸いです。この記事が役立つと思われた場合は、他の人も見つけられるように「いいね」を押して、お友達と共有してください。私の仕事の最新情報を入手するには、Linkedin で私をフォローしてください。読んでいただきありがとうございます!

以上がOpenVINO と Postgres を使用した高速かつ効率的なセマンティック検索システムの構築の詳細内容です。詳細については、PHP 中国語 Web サイトの他の関連記事を参照してください。