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 を検索システムに統合した方法は次のとおりです:

import argparse
import numpy as np
import pandas as pd
from typing import Any
from openvino.runtime import Core
from transformers import AutoTokenizer


def load_openvino_model(model_path: str) -> Core:
    core = Core()
    model = core.read_model(model_path   ".xml")
    compiled_model = core.compile_model(model, "CPU")
    return compiled_model


def rerank(
    compiled_model: Core,
    query: str,
    results: list[str],
    tokenizer: AutoTokenizer,
    batch_size: int,
) -> np.ndarray[np.float32, Any]:
    max_length = 512
    all_logits = []

    # Split results into batches
    for i in range(0, len(results), batch_size):
        batch_results = results[i : i   batch_size]
        inputs = tokenizer(
            [(query, item) for item in batch_results],
            padding=True,
            truncation="longest_first",
            max_length=max_length,
            return_tensors="np",
        )

        # Extract input tensors (convert to NumPy arrays)
        input_ids = inputs["input_ids"].astype(np.int32)
        attention_mask = inputs["attention_mask"].astype(np.int32)
        token_type_ids = inputs.get("token_type_ids", np.zeros_like(input_ids)).astype(
            np.int32
        )

        infer_request = compiled_model.create_infer_request()
        output = infer_request.infer(
            {
                "input_ids": input_ids,
                "attention_mask": attention_mask,
                "token_type_ids": token_type_ids,
            }
        )

        logits = output["logits"]
        all_logits.append(logits)

    all_logits = np.concatenate(all_logits, axis=0)
    return all_logits


def fetch_search_data(search_text: str) -> pd.DataFrame:
    # Usually you would fetch the data from a database
    df = pd.read_csv("cnbc_headlines.csv")
    df = df[~df["Headlines"].isnull()]

    texts = df["Headlines"].tolist()

    # Load the model and rerank
    openvino_model = load_openvino_model("cross-encoder-openvino-model/model")
    tokenizer = AutoTokenizer.from_pretrained("cross-encoder/ms-marco-TinyBERT-L-2-v2")
    rerank_scores = rerank(openvino_model, search_text, texts, tokenizer, batch_size=16)

    # Add the rerank scores to the DataFrame and sort by the new scores
    df["rerank_score"] = rerank_scores
    df = df.sort_values(by="rerank_score", ascending=False)

    return df


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Fetch search results with reranking using OpenVINO"
    )

    parser.add_argument(
        "--search_text",
        type=str,
        required=True,
        help="The search text to use for reranking",
    )

    args = parser.parse_args()

    df = fetch_search_data(args.search_text)
    print(df)

このアプローチにより、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 で私をフォローしてください。読んでいただきありがとうございます!