您的位置 首页 知识分享

使用 OpenVINO 和 Postgres 构建快速高效的语义搜索系统

照片由 real-napster 在 pixabay上 在我最近的一个项目中,我必须构建一个语义搜索系统,该系…

使用 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 分析了我的代码以可视化性能:

使用 OpenVINO 和 Postgres 构建快速高效的语义搜索系统

如您所见,标记化和预测步骤异常缓慢,导致搜索结果的提供出现严重延迟。总的来说,平均需要 4-5 秒。这是因为标记化和预测步骤之间存在阻塞操作。如果我们还添加其他操作,例如数据库调用、过滤等,我们很容易就总共需要 8-9 秒。

使用 openvino 优化性能

我面临的问题是:我们可以让它更快吗?答案是肯定的,通过利用 openvino,一个针对 cpu 推理优化的后端。 openvino 有助于加速英特尔硬件上的深度学习模型推理,我们在 aws lambda 上使用该硬件。

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 秒。完整的工作代码位于 hub 上。

速度微调:批量大小和标记化

提高性能的另一个关键因素是优化标记化流程并调整批量大小标记长度。通过增加批量大小(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中文网其它相关文章!

本文来自网络,不代表甲倪知识立场,转载请注明出处:http://www.spjiani.cn/wp/3714.html

作者: nijia

发表评论

您的电子邮箱地址不会被公开。

联系我们

联系我们

0898-88881688

在线咨询: QQ交谈

邮箱: email@wangzhan.com

工作时间:周一至周五,9:00-17:30,节假日休息

关注微信
微信扫一扫关注我们

微信扫一扫关注我们

关注微博
返回顶部