使用 OpenVINO 和 Postgres 建立快速且有效率的語意搜尋系統

照片由 real-napster 在 Pixabay上

在我最近的一個專案中，我必須建立一個語義搜尋系統，該系統可以高效能擴展並為報告搜尋提供即時回應。我們在 AWS RDS 上使用 PostgreSQL 和 pgvector，並搭配 AWS Lambda 來實現這一目標。面臨的挑戰是允許用戶使用自然語言查詢而不是依賴死板的關鍵字進行搜索，同時確保響應時間在 1-2 秒甚至更短，並且只能利用 CPU 資源。

在這篇文章中，我將逐步介紹建立此搜尋系統的步驟，從檢索到重新排名，以及使用 OpenVINO 和智慧批次進行標記化進行的最佳化。

語義搜尋概述：檢索和重新排序

現代最先進的搜尋系統通常包含兩個主要步驟：檢索和重新排名。

1) 檢索： 第一步涉及根據使用者查詢檢索相關文件的子集。這可以使用預先訓練的嵌入模型來完成，例如 OpenAI 的小型和大型嵌入、Cohere 的嵌入模型或 Mixbread 的 mxbai 嵌入。檢索的重點是透過測量文件與查詢的相似性來縮小文檔池的範圍。

這是一個使用 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 最佳化效能

我面臨的問題是：我們可以讓它更快嗎？ 答案是肯定的，透過利用 OpenVINO，一個針對 CPU 推理最佳化的後端。 OpenVINO 有助於加速英特爾硬體上的深度學習模型推理，我們在 AWS Lambda 上使用該硬體。

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 上。

速度微調：批量大小和標記化

提高效能的另一個關鍵因素是最佳化標記化流程並調整批次大小和標記長度。透過增加批次大小（batch_size = 16）和減少令牌長度（max_length = 512），我們可以並行化令牌化並減少重複操作的開銷。在我們的實驗中，我們發現 16 到 64 之間的 batch_size 效果很好，任何更大的值都會降低效能。同樣，我們將 max_length 設為 128，如果報告的平均長度相對較短，則該值是可行的。透過這些更改，我們實現了 8 倍的整體加速，將重新排名時間縮短至 1 秒以下，即使在 CPU 上也是如此。

在實踐中，這意味著嘗試不同的批量大小和令牌長度，以找到資料速度和準確性之間的適當平衡。透過這樣做，我們看到回應時間顯著縮短，使得搜尋系統即使有 1,000 份報告也可擴展。

結論

透過使用 OpenVINO 並最佳化標記化和批次處理，我們能夠建立一個高效能語意搜尋系統，滿足僅 CPU 設定的即時要求。事實上，我們的整體速度提升了 8 倍。使用句子轉換器進行檢索與使用交叉編碼器模型進行重新排名相結合，創造了強大的、用戶友好的搜尋體驗。

如果您正在建立回應時間和運算資源受到限制的類似系統，我強烈建議您探索 OpenVINO 和智慧批次以釋放更好的效能。

希望您喜歡這篇文章。如果您覺得這篇文章有用，請給我一個贊，以便其他人也可以找到它，並與您的朋友分享。在 Linkedin 上關注我，以了解我的最新工作。感謝您的閱讀！

以上是使用 OpenVINO 和 Postgres 建立快速且有效率的語意搜尋系統的詳細內容。更多資訊請關注PHP中文網其他相關文章！