使用 HuggingFace 库进行本地嵌入向量生成

时间：2025-03-24 10:20:56

在当今的AI和机器学习应用中，嵌入向量（embeddings）已成为不可或缺的一部分。嵌入向量能够将文本等高维数据转换为低维稠密向量，从而便于计算和分析。在本文中，我们将介绍如何使用 HuggingFace 库在本地生成嵌入向量，并演示相关代码。

环境准备

首先，我们需要安装一些必要的依赖库。可以通过以下命令进行安装：

# 安装必要的库
!pip install sentence-transformers
!pip install llama-index
!pip install optimum[exporters]
!pip install InstructorEmbedding
!pip install -U sentence-transformers==2.2.2

使用 HuggingFace 嵌入模型

HuggingFace 提供了多种嵌入模型。以下是使用基础 HuggingFace 嵌入模型的代码示例：

from llama_index.embeddings.huggingface import HuggingFaceEmbedding

# 加载 BAAI/bge-small-en-v1.5 模型
embed_model = HuggingFaceEmbedding(model_name="BAAI/bge-small-en-v1.5")

# 获取文本的嵌入向量
embeddings = embed_model.get_text_embedding("Hello World!")
print(len(embeddings))
print(embeddings[:5])

使用指令嵌入模型

如果希望嵌入向量能够根据指令进行调整，可以使用指令嵌入模型：

from llama_index.embeddings.instructor import InstructorEmbedding

embed_model = InstructorEmbedding(model_name="hkunlp/instructor-base")

# 获取文本的嵌入向量
embeddings = embed_model.get_text_embedding("Hello World!")
print(len(embeddings))
print(embeddings[:5])

使用 Optimum 库创建 ONNX 模型

Optimum 是一个用于导出和运行 HuggingFace 模型的库，支持 ONNX 格式。以下是创建和使用 ONNX 模型的示例代码：

from llama_index.embeddings.huggingface_optimum import OptimumEmbedding

# 创建并保存 ONNX 模型
OptimumEmbedding.create_and_save_optimum_model("BAAI/bge-small-en-v1.5", "./bge_onnx")

# 使用创建的 ONNX 模型
embed_model = OptimumEmbedding(folder_name="./bge_onnx")

# 获取文本的嵌入向量
embeddings = embed_model.get_text_embedding("Hello World!")
print(len(embeddings))
print(embeddings[:5])

性能比较

为了验证模型的性能，我们可以下载一个大型文档（如 IPCC 气候报告的第三章），并对其进行嵌入向量生成：

# 下载 IPCC 气候报告第三章
!curl https://www.ipcc.ch/report/ar6/wg2/downloads/report/IPCC_AR6_WGII_Chapter03.pdf --output IPCC_AR6_WGII_Chapter03.pdf

from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
from llama_index.core import Settings

# 读取并加载文档
documents = SimpleDirectoryReader(input_files=["IPCC_AR6_WGII_Chapter03.pdf"]).load_data()

from llama_index.embeddings.huggingface import HuggingFaceEmbedding

# 使用 HuggingFace 嵌入模型
embed_model = HuggingFaceEmbedding(model_name="BAAI/bge-small-en-v1.5")
Settings.embed_model = embed_model

%timeit -r 1 -n 1
index = VectorStoreIndex.from_documents(documents, show_progress=True)

from llama_index.embeddings.huggingface_optimum import OptimumEmbedding

# 使用 Optimum 嵌入模型
embed_model = OptimumEmbedding(folder_name="./bge_onnx")
Settings.embed_model = embed_model

%timeit -r 1 -n 1
index = VectorStoreIndex.from_documents(documents, show_progress=True)

可能遇到的错误

  1. Can’t initialize NVML 警告：如果在运行代码时看到 Can't initialize NVML 警告，可以忽略它，因为这通常与 GPU 驱动程序有关，不会影响嵌入向量的生成。
  2. 环境变量未设置：在使用 OpenAI API 时，请确保设置了 OPENAI_API_KEY 环境变量，否则会导致 API 调用失败。

参考资料:
HuggingFace,LlamaIndex。

如果你觉得这篇文章对你有帮助，请点赞，关注我的博客，谢谢!
title: 使用 HuggingFace 库进行本地嵌入向量生成
body: 在当今的AI和机器学习应用中，嵌入向量（embeddings）已成为不可或缺的一部分。嵌入向量能够将文本等高维数据转换为低维稠密向量，从而便于计算和分析。在本文中，我们将介绍如何使用 HuggingFace 库在本地生成嵌入向量，并演示相关代码。

环境准备

首先，我们需要安装一些必要的依赖库。可以通过以下命令进行安装：

# 安装必要的库
!pip install sentence-transformers
!pip install llama-index
!pip install optimum[exporters]
!pip install InstructorEmbedding
!pip install -U sentence-transformers==2.2.2

使用 HuggingFace 嵌入模型

HuggingFace 提供了多种嵌入模型。以下是使用基础 HuggingFace 嵌入模型的代码示例：

from llama_index.embeddings.huggingface import HuggingFaceEmbedding

# 加载 BAAI/bge-small-en-v1.5 模型
embed_model = HuggingFaceEmbedding(model_name="BAAI/bge-small-en-v1.5")

# 获取文本的嵌入向量
embeddings = embed_model.get_text_embedding("Hello World!")
print(len(embeddings))
print(embeddings[:5])

使用指令嵌入模型

如果希望嵌入向量能够根据指令进行调整，可以使用指令嵌入模型：

from llama_index.embeddings.instructor import InstructorEmbedding

embed_model = InstructorEmbedding(model_name="hkunlp/instructor-base")

# 获取文本的嵌入向量
embeddings = embed_model.get_text_embedding("Hello World!")
print(len(embeddings))
print(embeddings[:5])

使用 Optimum 库创建 ONNX 模型

Optimum 是一个用于导出和运行 HuggingFace 模型的库，支持 ONNX 格式。以下是创建和使用 ONNX 模型的示例代码：

from llama_index.embeddings.huggingface_optimum import OptimumEmbedding

# 创建并保存 ONNX 模型
OptimumEmbedding.create_and_save_optimum_model("BAAI/bge-small-en-v1.5", "./bge_onnx")

# 使用创建的 ONNX 模型
embed_model = OptimumEmbedding(folder_name="./bge_onnx")

# 获取文本的嵌入向量
embeddings = embed_model.get_text_embedding("Hello World!")
print(len(embeddings))
print(embeddings[:5])

性能比较

为了验证模型的性能，我们可以下载一个大型文档（如 IPCC 气候报告的第三章），并对其进行嵌入向量生成：

# 下载 IPCC 气候报告第三章
!curl https://www.ipcc.ch/report/ar6/wg2/downloads/report/IPCC_AR6_WGII_Chapter03.pdf --output IPCC_AR6_WGII_Chapter03.pdf

from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
from llama_index.core import Settings

# 读取并加载文档
documents = SimpleDirectoryReader(input_files=["IPCC_AR6_WGII_Chapter03.pdf"]).load_data()

from llama_index.embeddings.huggingface import HuggingFaceEmbedding

# 使用 HuggingFace 嵌入模型
embed_model = HuggingFaceEmbedding(model_name="BAAI/bge-small-en-v1.5")
Settings.embed_model = embed_model

%timeit -r 1 -n 1
index = VectorStoreIndex.from_documents(documents, show_progress=True)

from llama_index.embeddings.huggingface_optimum import OptimumEmbedding

# 使用 Optimum 嵌入模型
embed_model = OptimumEmbedding(folder_name="./bge_onnx")
Settings.embed_model = embed_model

%timeit -r 1 -n 1
index = VectorStoreIndex.from_documents(documents, show_progress=True)

可能遇到的错误

  1. Can’t initialize NVML 警告：如果在运行代码时看到 Can't initialize NVML 警告，可以忽略它，因为这通常与 GPU 驱动程序有关，不会影响嵌入向量的生成。
  2. 环境变量未设置：在使用 OpenAI API 时，请确保设置了 OPENAI_API_KEY 环境变量，否则会导致 API 调用失败。

参考资料:
HuggingFace,LlamaIndex。

如果你觉得这篇文章对你有帮助，请点赞，关注我的博客，谢谢!

相关文章