4.1 检索后处理

LlamaIndex 的 Node Postprocessors 提供了一系列检索后处理模块，用于在 Node 检索步骤之后、响应合成步骤之前，对检索到的 Node 进行转换或过滤。

例如：我们可以用不同模型对检索后的 Nodes 做重排序：

import chromadb
from chromadb import Settings
from llama_index.core import VectorStoreIndex
from llama_index.vector_stores.chroma import ChromaVectorStore
from llama_index.core import StorageContext
from llama_index.core import SimpleDirectoryReader
from llama_index.readers.file import PyMuPDFReader
from llama_index.core.node_parser import TokenTextSplitter
from llama_index.core.postprocessor import SentenceTransformerRerank

reader = SimpleDirectoryReader(
    input_dir="./data",  # 目标目录
    recursive=False,  # 是否递归遍历子目录
    required_exts=[".pdf"],  # (可选)只读取指定后缀的文件
    file_extractor={".pdf": PyMuPDFReader()}  # (可选)指定文件读取器
)

# 读取并解析pdf
documents = reader.load_data()

node_parser = TokenTextSplitter(chunk_size=300, chunk_overlap=100)

nodes = node_parser.get_nodes_from_documents(
    documents, show_progress=False
)

# 创建 ChromaDB 实例
chroma_client = chromadb.HttpClient(settings=Settings(allow_reset=True), port=8000)
chroma_client.reset()
chroma_collection = chroma_client.create_collection("llama_index_demo")

# 创建 ChromaDB VectorStore
vector_store = ChromaVectorStore(chroma_collection=chroma_collection)
storage_context = StorageContext.from_defaults(vector_store=vector_store)
chroma_index = VectorStoreIndex(nodes, storage_context=storage_context)

# 获取 retriever
vector_retriever = chroma_index.as_retriever(similarity_top_k=5)

# 检索
query_str = "Llama 2使用的是transformer吗"
results = vector_retriever.retrieve(query_str)
print("------------------------排序前--------------------------")
for i, node in enumerate(results):
    print(f"[{i}] {node.text}")
    print()

# 检索后排序模型
postprocessor = SentenceTransformerRerank(
    model="BAAI/bge-reranker-large", top_n=2
)
print("------------------------排序后--------------------------")
sort_nodes = postprocessor.postprocess_nodes(results, query_str=query_str)

for i, node in enumerate(sort_nodes):
    print(f"[{i}] {node.text}")
    print()

可以看到排序后 transformer 相关的内容被提前，重排序模型能够更精准地识别与问题相关的内容。

4.2 对话

import chromadb
from chromadb import Settings
from llama_index.core import VectorStoreIndex
from llama_index.vector_stores.chroma import ChromaVectorStore


# 创建 ChromaDB 实例
chroma_client = chromadb.HttpClient(settings=Settings(allow_reset=True), port=8000)
chroma_collection = chroma_client.get_or_create_collection("llama_index_demo")

# 创建 ChromaDB VectorStore
vector_store = ChromaVectorStore(chroma_collection=chroma_collection)
chroma_index = VectorStoreIndex.from_vector_store(vector_store)

# 单轮问答
print("单轮问答:")
qa_engine = chroma_index.as_query_engine()
response = qa_engine.query("Llama2 有多少参数?")
print(response)

# 流式输出
print("流式输出:")
qa_engine = chroma_index.as_query_engine(streaming=True)
response = qa_engine.query("Llama2 有多少参数?")
response.print_response_stream()
print()

# 多轮对话
print("多轮对话:")
chat_engine = chroma_index.as_chat_engine()
response = chat_engine.chat("Llama2 有多少参数?")
print(response)
response = chat_engine.chat("最多有多少？")
print(response)