RAG（4）：自动合并检索Auto-merging Retrieval

• 自动合并检索的使用

  from llama_index import Document
  document = Document(text="\n\n".join([doc.text for doc in documents]))
  from llama_index.node_parser import HierarchicalNodeParser
  
  ## 创建parser
  # create the hierarchical node parser w/ default settings
  node_parser = HierarchicalNodeParser.from_defaults(
      chunk_sizes=[2048, 512, 128]
  )
  nodes = node_parser.get_nodes_from_documents([document])
  
  ## 只检索叶节点
  from llama_index.node_parser import get_leaf_nodes
  leaf_nodes = get_leaf_nodes(nodes)
  print(leaf_nodes[30].text)
  nodes_by_id = {node.node_id: node for node in nodes}
  parent_node = nodes_by_id[leaf_nodes[30].parent_node.node_id]
  print(parent_node.text)
  
  ## 使用BGE嵌入模型
  from llama_index.llms import OpenAI
  llm = OpenAI(model="gpt-3.5-turbo", temperature=0.1)
  from llama_index import ServiceContext
  auto_merging_context = ServiceContext.from_defaults(
      llm=llm,
      embed_model="local:BAAI/bge-small-en-v1.5",
      node_parser=node_parser,
  )
  
  ## 存储上下文对象
  from llama_index import VectorStoreIndex, StorageContext
  storage_context = StorageContext.from_defaults()
  storage_context.docstore.add_documents(nodes)
  automerging_index = VectorStoreIndex(
      leaf_nodes, storage_context=storage_context, service_context=auto_merging_context
  )
  automerging_index.storage_context.persist(persist_dir="./merging_index")
  
  ## 设置检索器并运行Query
  # This block of code is optional to check
  # if an index file exist, then it will load it
  # if not, it will rebuild it
  import os
  from llama_index import VectorStoreIndex, StorageContext, load_index_from_storage
  from llama_index import load_index_from_storage
  if not os.path.exists("./merging_index"):
      storage_context = StorageContext.from_defaults()
      storage_context.docstore.add_documents(nodes)
  
      automerging_index = VectorStoreIndex(
              leaf_nodes,
              storage_context=storage_context,
              service_context=auto_merging_context
          )
  
      automerging_index.storage_context.persist(persist_dir="./merging_index")
  else:
      automerging_index = load_index_from_storage(
          StorageContext.from_defaults(persist_dir="./merging_index"),
          service_context=auto_merging_context
      )
  
  

  import os
  
  from llama_index import (
      ServiceContext,
      StorageContext,
      VectorStoreIndex,
      load_index_from_storage,
  )
  from llama_index.node_parser import HierarchicalNodeParser
  from llama_index.node_parser import get_leaf_nodes
  from llama_index import StorageContext, load_index_from_storage
  from llama_index.retrievers import AutoMergingRetriever
  from llama_index.indices.postprocessor import SentenceTransformerRerank
  from llama_index.query_engine import RetrieverQueryEngine
  
  def build_automerging_index(
      documents,
      llm,
      embed_model="local:BAAI/bge-small-en-v1.5",
      save_dir="merging_index",
      chunk_sizes=None,
  ):
      chunk_sizes = chunk_sizes or [2048, 512, 128]
      node_parser = HierarchicalNodeParser.from_defaults(chunk_sizes=chunk_sizes)
      nodes = node_parser.get_nodes_from_documents(documents)
      leaf_nodes = get_leaf_nodes(nodes)
      merging_context = ServiceContext.from_defaults(
          llm=llm,
          embed_model=embed_model,
      )
      storage_context = StorageContext.from_defaults()
      storage_context.docstore.add_documents(nodes)
  
      if not os.path.exists(save_dir):
          automerging_index = VectorStoreIndex(
              leaf_nodes, storage_context=storage_context, service_context=merging_context
          )
          automerging_index.storage_context.persist(persist_dir=save_dir)
      else:
          automerging_index = load_index_from_storage(
              StorageContext.from_defaults(persist_dir=save_dir),
              service_context=merging_context,
          )
      return automerging_index
  
  def get_automerging_query_engine(
      automerging_index,
      similarity_top_k=12,
      rerank_top_n=6,
  ):
      base_retriever = automerging_index.as_retriever(similarity_top_k=similarity_top_k)
      retriever = AutoMergingRetriever(
          base_retriever, automerging_index.storage_context, verbose=True
      )
      rerank = SentenceTransformerRerank(
          top_n=rerank_top_n, model="BAAI/bge-reranker-base"
      )
      auto_merging_engine = RetrieverQueryEngine.from_args(
          retriever, node_postprocessors=[rerank]
      )
      return auto_merging_engine
  

• 评估自动合并检索器

  ## 设置index，两层（不固定）
  auto_merging_index_0 = build_automerging_index(
      documents,
      llm=OpenAI(model="gpt-3.5-turbo", temperature=0.1),
      embed_model="local:BAAI/bge-small-en-v1.5",
      save_dir="merging_index_0",
      chunk_sizes=[2048,512],
  )
  ## 设置自动合并器
  auto_merging_engine_0 = get_automerging_query_engine(
      auto_merging_index_0,
      similarity_top_k=12,
      rerank_top_n=6,
  )
  from utils import get_prebuilt_trulens_recorder
  tru_recorder = get_prebuilt_trulens_recorder(
      auto_merging_engine_0,
      app_id ='app_0'
  )
  ## 加载用于评估的问题
  eval_questions = []
  with open('generated_questions.text', 'r') as file:
      for line in file:
          # Remove newline character and convert to integer
          item = line.strip()
          eval_questions.append(item)
  ## 评估对应问答
  def run_evals(eval_questions, tru_recorder, query_engine):
      for question in eval_questions:
          with tru_recorder as recording:
              response = query_engine.query(question)
  run_evals(eval_questions, tru_recorder, auto_merging_engine_0)
  ## 运行
  from trulens_eval import Tru
  Tru().get_leaderboard(app_ids=[])
  

  ## 设置index 三层
  auto_merging_index_1 = build_automerging_index(
      documents,
      llm=OpenAI(model="gpt-3.5-turbo", temperature=0.1),
      embed_model="local:BAAI/bge-small-en-v1.5",
      save_dir="merging_index_1",
      chunk_sizes=[2048,512,128],
  )
  
  ##
  auto_merging_engine_1 = get_automerging_query_engine(
      auto_merging_index_1,
      similarity_top_k=12,
      rerank_top_n=6,
  )
  
  ##
  tru_recorder = get_prebuilt_trulens_recorder(
      auto_merging_engine_1,
      app_id ='app_1'
  )
  
  ## 
  run_evals(eval_questions, tru_recorder, auto_merging_engine_1)
  
  ##
  from trulens_eval import Tru
  Tru().get_leaderboard(app_ids=[])