设置¶
如果您在colab上打开这个笔记本,您可能需要安装LlamaIndex 🦙。
%pip install llama-index-llms-openai
%pip install llama-index-retrievers-bm25
!pip install llama-index
# 注意:这仅在jupyter笔记本中是必要的。
# 详情:Jupyter在后台运行一个事件循环。
# 当我们启动一个事件循环来进行异步查询时,会导致嵌套的事件循环。
# 通常情况下是不允许的,我们使用nest_asyncio来允许这种情况以方便操作。
import nest_asyncio
nest_asyncio.apply()
import os
import openai
os.environ["OPENAI_API_KEY"] = "sk-..."
openai.api_key = os.environ["OPENAI_API_KEY"]
import logging
import sys
logging.basicConfig(stream=sys.stdout, level=logging.INFO)
logging.getLogger().handlers = []
logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))
from llama_index.core import (
SimpleDirectoryReader,
StorageContext,
VectorStoreIndex,
)
from llama_index.retrievers.bm25 import BM25Retriever
from llama_index.core.retrievers import VectorIndexRetriever
from llama_index.core.node_parser import SentenceSplitter
from llama_index.llms.openai import OpenAI
下载数据¶
!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()
# 初始化LLM + 节点解析器
llm = OpenAI(model="gpt-4")
splitter = SentenceSplitter(chunk_size=1024)
nodes = splitter.get_nodes_from_documents(documents)
# 初始化存储上下文(默认情况下是内存中的)
storage_context = StorageContext.from_defaults()
storage_context.docstore.add_documents(nodes)
index = VectorStoreIndex(
nodes=nodes,
storage_context=storage_context,
)
HTTP Request: POST https://api.openai.com/v1/embeddings "HTTP/1.1 200 OK"
BM25检索器¶
我们将使用BM25检索器来搜索文档。
# 我们可以传入索引、doctore或节点列表来创建检索器
retriever = BM25Retriever.from_defaults(nodes=nodes, similarity_top_k=2)
from llama_index.core.response.notebook_utils import display_source_node
# 将从特定公司检索上下文
nodes = retriever.retrieve("Viaweb和Interleaf发生了什么?")
for node in nodes:
display_source_node(node)
Node ID: 9afb8cb9-42f3-4160-807c-1cd6685fa774
Similarity: 1.6781810094192822
Text: Now that I could write essays again, I wrote a bunch about topics I'd had stacked up. I kept writ...
Node ID: 71de4371-10ff-4d3a-8a31-14b2d4ddec83
Similarity: 1.546534805781164
Text: I couldn't have put this into words when I was 18. All I knew at the time was that I kept taking ...
nodes = retriever.retrieve("What did Paul Graham do after RISD?")
for node in nodes:
display_source_node(node)
Node ID: 42c88008-dd05-4590-8fd1-df85567579ea
Similarity: 5.389397745654172
Text: It was missing a lot of things you'd want in a programming language. So these had to be added, an...
Node ID: 4d98c2ad-60cc-4cc0-abe2-b95c0c4be87b
Similarity: 1.141523170922594
Text: Painting students were supposed to express themselves, which to the more worldly ones meant to tr...
使用BM25方法的路由器检索器¶
现在我们将结合BM25检索器和向量索引检索器。
from llama_index.core.tools import RetrieverTool
vector_retriever = VectorIndexRetriever(index)
bm25_retriever = BM25Retriever.from_defaults(nodes=nodes, similarity_top_k=2)
retriever_tools = [
RetrieverTool.from_defaults(
retriever=vector_retriever,
description="Useful in most cases",
),
RetrieverTool.from_defaults(
retriever=bm25_retriever,
description="Useful if searching about specific information",
),
]
from llama_index.core.retrievers import RouterRetriever
retriever = RouterRetriever.from_defaults(
retriever_tools=retriever_tools,
llm=llm,
select_multi=True,
)
# 将从作者的生活中检索所有上下文
nodes = retriever.retrieve(
"您能给我关于作者生活的所有上下文吗?"
)
for node in nodes:
display_source_node(node)
HTTP Request: POST https://api.openai.com/v1/chat/completions "HTTP/1.1 200 OK" Selecting retriever 0: The author's life context is a broad topic and can be useful in most cases.. HTTP Request: POST https://api.openai.com/v1/embeddings "HTTP/1.1 200 OK"
Node ID: 9afb8cb9-42f3-4160-807c-1cd6685fa774
Similarity: 0.7961776744788842
Text: Now that I could write essays again, I wrote a bunch about topics I'd had stacked up. I kept writ...
Node ID: b1ffb78d-dc4c-439b-906a-cca73b713940
Similarity: 0.7924813308773564
Text: We actually had one of those little stoves, fed with kindling, that you see in 19th century studi...
高级 - 混合检索器 + 重新排序¶
在这里,我们扩展基本的检索器类,并创建一个自定义的检索器,它始终使用向量检索器和BM25检索器。
然后,节点可以被重新排序和过滤。这样我们可以保持中间的top-k值很大,并让重新排序过滤掉不需要的节点。
为了最好地演示这一点,我们将使用一个更大的源文档集合——来自2022年IPCC气候报告的第3章。
设置数据¶
!curl https://www.ipcc.ch/report/ar6/wg2/downloads/report/IPCC_AR6_WGII_Chapter03.pdf --output IPCC_AR6_WGII_Chapter03.pdf
# !pip install pypdf
from llama_index.core import VectorStoreIndex,StorageContext,SimpleDirectoryReader和Document
from llama_index.core.node_parser import SentenceSplitter
from llama_index.llms.openai import OpenAI
# 加载文档
documents = SimpleDirectoryReader(
input_files=["IPCC_AR6_WGII_Chapter03.pdf"]
).load_data()
# 初始化llm + 节点解析器
# -- 在这里,我们设置了一个较小的块大小,以便更有效地重新排名
llm = OpenAI(model="gpt-3.5-turbo")
splitter = SentenceSplitter(chunk_size=256)
# 限制到一个较小的部分
nodes = splitter.get_nodes_from_documents(
[Document(text=documents[0].get_content()[:1000000])]
)
# 初始化存储上下文(默认情况下是内存中的)
storage_context = StorageContext.from_defaults()
storage_context.docstore.add_documents(nodes)
index = VectorStoreIndex(nodes, storage_context=storage_context)
from llama_index.retrievers.bm25 import BM25Retriever
# 使用嵌入检索出前10个最相似的节点
vector_retriever = index.as_retriever(similarity_top_k=10)
# 使用bm25检索出前10个最相似的节点
bm25_retriever = BM25Retriever.from_defaults(nodes=nodes, similarity_top_k=10)
class CustomRetriever:
def __init__(self, data):
self.data = data
def retrieve(self, key):
if key in self.data:
return self.data[key]
else:
return "Key not found"
自定义检索器实现
from llama_index.core.retrievers import BaseRetriever
class HybridRetriever(BaseRetriever):
def __init__(self, vector_retriever, bm25_retriever):
self.vector_retriever = vector_retriever
self.bm25_retriever = bm25_retriever
super().__init__()
def _retrieve(self, query, **kwargs):
bm25_nodes = self.bm25_retriever.retrieve(query, **kwargs)
vector_nodes = self.vector_retriever.retrieve(query, **kwargs)
# 合并两个节点列表
all_nodes = []
node_ids = set()
for n in bm25_nodes + vector_nodes:
if n.node.node_id not in node_ids:
all_nodes.append(n)
node_ids.add(n.node.node_id)
return all_nodes
index.as_retriever(similarity_top_k=5)
hybrid_retriever = HybridRetriever(vector_retriever, bm25_retriever)
重新排序器设置¶
!pip install sentence-transformers
from llama_index.core.postprocessor import SentenceTransformerRerank
reranker = SentenceTransformerRerank(top_n=4, model="BAAI/bge-reranker-base")
检索¶
from llama_index.core import QueryBundle
from llama_index.core.schema import NodeWithScore
retrieved_nodes = hybrid_retriever.retrieve(
"What is the impact of climate change on the ocean?"
)
reranked_nodes = reranker.postprocess_nodes(
retrieved_nodes,
query_bundle=QueryBundle(
"What is the impact of climate change on the ocean?"
),
)
print("Initial retrieval: ", len(retrieved_nodes), " nodes")
print("Re-ranked retrieval: ", len(reranked_nodes), " nodes")
from llama_index.core.response.notebook_utils import display_source_node
for node in reranked_nodes:
display_source_node(node)
Node ID: 735c563d-e68c-42e7-bbeb-d360a2918d22
Similarity: 0.6910567283630371
Text: lb88\8a|hiac2:,sg
`HfvyȔskz4
4+)t$EAs9<"L䴺-iai]}?R)Jv8Q0V#9>f=3ב&WTFKNӅ9GN8v
4g
tz...
Node ID: 07a796db-8307-47a0-83a0-9b5f953b86b0
Similarity: 0.609274685382843
Text: {_d3<,fCȀ
0K0~n
(A ̙$saxt
:
xY3z/ߕXAWwpTeHY0HZHe̚kÇ>.82%ϖ^.rCS.2^12C?
Node ID: 0ee98fbd-d015-4cd4-aba0-812423880915
Similarity: 0.5765283107757568
Text: X+V/Z
6;O8/%Z3EC\asU1f(xLͼ]X\q"1}.66OKGSi9ǧ'(?Iʲi=4ޮ^m,Mp1~lBxP3]S<?C)3WM;WGQt@l...
Node ID: 2562bbfb-fda2-432d-9674-5ba77ff4b767
Similarity: 0.3882860541343689
Text: PxS0s.MAK|+MOUÊ^;tc0
$XHmL<ҁ;Q@nUd&ի
2`B01i^yOw:rsM冫S*wu4a{-.Խ=j鷩nQD_AA%~T>~'OxU gk...
完整查询引擎¶
from llama_index.core.query_engine import RetrieverQueryEngine
query_engine = RetrieverQueryEngine.from_args(
retriever=hybrid_retriever,
node_postprocessors=[reranker],
llm=llm,
)
response = query_engine.query(
"What is the impact of climate change on the ocean?"
)
from llama_index.core.response.notebook_utils import display_response
display_response(response)
Final Response:
Climate change can have a significant impact on the ocean. Rising temperatures can lead to the melting of polar ice caps, resulting in sea-level rise and coastal flooding. It can also disrupt ocean currents and affect marine ecosystems, including coral reefs and fish populations. Additionally, climate change can lead to ocean acidification, which can harm marine life and impact the overall health of the ocean.