利用RAG（检索增强生成）技术实现多文档检索

最新推荐文章于 2024-09-09 00:00:00 发布

qq_37836323

最新推荐文章于 2024-09-09 00:00:00 发布

阅读量340

点赞数 5

文章标签： python

本文链接：https://blog.csdn.net/qq_29929123/article/details/140316675

版权

在处理多文档检索时，使用RAG（检索增强生成）框架能帮助我们更有效地从不同文档中获取相关信息。本文将介绍如何使用LlamaIndex和Weaviate来实现这种多文档检索。

安装依赖

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

%pip install llama-index-readers-github
%pip install llama-index-vector-stores-weaviate
!pip install llama-index llama-hub

下载和设置数据

在此部分，我们将加载LlamaIndex的GitHub问题数据。

import nest_asyncio
import os
from llama_index.readers.github import GitHubRepositoryIssuesReader, GitHubIssuesClient

nest_asyncio.apply()

os.environ["GITHUB_TOKEN"] = "ghp_..."
os.environ["OPENAI_API_KEY"] = "sk-..."

github_client = GitHubIssuesClient()
loader = GitHubRepositoryIssuesReader(github_client, owner="run-llama", repo="llama_index", verbose=True)
orig_docs = loader.load_data()

# 限制为100个文档
limit = 100
docs = []
for idx, doc in enumerate(orig_docs):
    doc.metadata["index_id"] = int(doc.id_)
    if idx >= limit:
        break
    docs.append(doc)

设置向量存储和索引

import weaviate
from llama_index.vector_stores.weaviate import WeaviateVectorStore
from llama_index.core import VectorStoreIndex, StorageContext

auth_config = weaviate.AuthApiKey(api_key="XRa15cDIkYRT7AkrpqT6jLfE4wropK1c1TGk")
client = weaviate.Client("https://llama-index-test-v0oggsoz.weaviate.network", auth_client_secret=auth_config)
class_name = "LlamaIndex_docs"

client.schema.delete_class(class_name)  # 可选：删除已有的Schema

vector_store = WeaviateVectorStore(weaviate_client=client, index_name=class_name)
storage_context = StorageContext.from_defaults(vector_store=vector_store)
doc_index = VectorStoreIndex.from_documents(docs, storage_context=storage_context)

创建检索和过滤的IndexNode

from llama_index.core import SummaryIndex
from llama_index.core.async_utils import run_jobs
from llama_index.llms.openai import OpenAI
from llama_index.core.schema import IndexNode
from llama_index.core.vector_stores import FilterOperator, MetadataFilter, MetadataFilters

async def aprocess_doc(doc, include_summary: bool = True):
    metadata = doc.metadata
    date_tokens = metadata["created_at"].split("T")[0].split("-")
    year, month, day = int(date_tokens[0]), int(date_tokens[1]), int(date_tokens[2])
    assignee = "" if "assignee" not in doc.metadata else doc.metadata["assignee"]
    size = ""
    if len(doc.metadata["labels"]) > 0:
        size_arr = [l for l in doc.metadata["labels"] if "size:" in l]
        size = size_arr[0].split(":")[1] if len(size_arr) > 0 else ""
    new_metadata = {"state": metadata["state"], "year": year, "month": month, "day": day, "assignee": assignee, "size": size}

    summary_index = SummaryIndex.from_documents([doc])
    query_str = "Give a one-sentence concise summary of this issue."
    query_engine = summary_index.as_query_engine(llm=OpenAI(model="gpt-3.5-turbo"))
    summary_txt = await query_engine.aquery(query_str)
    summary_txt = str(summary_txt)

    index_id = doc.metadata["index_id"]
    filters = MetadataFilters(filters=[MetadataFilter(key="index_id", operator=FilterOperator.EQ, value=int(index_id))])

    index_node = IndexNode(text=summary_txt, metadata=new_metadata, obj=doc_index.as_retriever(filters=filters), index_id=doc.id_)
    return index_node

async def aprocess_docs(docs):
    tasks = [aprocess_doc(doc) for doc in docs]
    index_nodes = await run_jobs(tasks, show_progress=True, workers=3)
    return index_nodes

index_nodes = await aprocess_docs(docs)

创建顶级AutoRetriever

加载总结的元数据和原始文档到向量数据库中，可以执行结构化、分层的检索策略。

import weaviate
from llama_index.vector_stores.weaviate import WeaviateVectorStore
from llama_index.core import VectorStoreIndex, StorageContext

auth_config = weaviate.AuthApiKey(api_key="XRa15cDIkYRT7AkrpqT6jLfE4wropK1c1TGk")
client = weaviate.Client("https://llama-index-test-v0oggsoz.weaviate.network", auth_client_secret=auth_config)
class_name = "LlamaIndex_auto"

client.schema.delete_class(class_name)  # 可选：删除已有的Schema

vector_store_auto = WeaviateVectorStore(weaviate_client=client, index_name=class_name)
storage_context_auto = StorageContext.from_defaults(vector_store=vector_store_auto)

index = VectorStoreIndex(objects=index_nodes, storage_context=storage_context_auto)

设置可组合的Auto-Retriever

from llama_index.core.vector_stores import MetadataInfo, VectorStoreInfo

vector_store_info = VectorStoreInfo(
    content_info="Github Issues",
    metadata_info=[
        MetadataInfo(name="state", description="Whether the issue is `open` or `closed`", type="string"),
        MetadataInfo(name="year", description="The year issue was created", type="integer"),
        MetadataInfo(name="month", description="The month issue was created", type="integer"),
        MetadataInfo(name="day", description="The day issue was created", type="integer"),
        MetadataInfo(name="assignee", description="The assignee of the ticket", type="string"),
        MetadataInfo(name="size", description="How big the issue is (XS, S, M, L, XL, XXL)", type="string"),
    ],
)

from llama_index.core.retrievers import VectorIndexAutoRetriever

retriever = VectorIndexAutoRetriever(
    index,
    vector_store_info=vector_store_info,
    similarity_top_k=2,
    empty_query_top_k=10,
    verbose=True,
)

尝试检索

from llama_index.core import QueryBundle

nodes = retriever.retrieve(QueryBundle("Tell me about some issues on 01/11"))
print(f"Number of source nodes: {len(nodes)}")
print(nodes[0].node.metadata)

集成RetrieverQueryEngine

from llama_index.core.query_engine import RetrieverQueryEngine
from llama_index.llms.openai import OpenAI

llm = OpenAI(model="gpt-3.5-turbo")
query_engine = RetrieverQueryEngine.from_args(retriever, llm=llm)

response = query_engine.query("Tell me about some issues on 01/11")
print(str(response))