RAG（3）：句子窗口检索Sentence Window Retrieval

• 使用句窗检索

  句窗检索机制同基本RAG对比

  

  

  # 设置句窗检索
  ## 1.创建句窗节点解析器
  from llama_index.node_parser import SentenceWindowNodeParser
  ## 将文档分割为句子，并利用句子周围的上下文增强句子块
  # create the sentence window node parser w/ default settings
  node_parser = SentenceWindowNodeParser.from_defaults(
      window_size=3,
      window_metadata_key="window",
      original_text_metadata_key="original_text",
  )
  
  ## 2.llm接口设置
  from llama_index.llms import OpenAI
  llm = OpenAI(model="gpt-3.5-turbo", temperature=0.1)
  
  ## 3.服务上下文对象设置
  from llama_index import ServiceContext
  sentence_context = ServiceContext.from_defaults(
      llm=llm,
      embed_model="local:BAAI/bge-small-en-v1.5",
      # embed_model="local:BAAI/bge-large-en-v1.5"
      node_parser=node_parser,
  )
  
  ## 4.设置包含源文档的向量存储索引
  from llama_index import VectorStoreIndex
  sentence_index = VectorStoreIndex.from_documents(
      [document], service_context=sentence_context
  )
  sentence_index.storage_context.persist(persist_dir="./sentence_index")
  # 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("./sentence_index"):
      sentence_index = VectorStoreIndex.from_documents(
          [document], service_context=sentence_context
      )
  
      sentence_index.storage_context.persist(persist_dir="./sentence_index")
  else:
      sentence_index = load_index_from_storage(
          StorageContext.from_defaults(persist_dir="./sentence_index"),
          service_context=sentence_context
      )
      
  ## 5.元数据替换后处理器
  from llama_index.schema import NodeWithScore
  from copy import deepcopy
  scored_nodes = [NodeWithScore(node=x, score=1.0) for x in nodes]
  nodes_old = [deepcopy(n) for n in nodes]
  ## 基于BGE嵌入的重新排序器
  from llama_index.indices.postprocessor import SentenceTransformerRerank
  # BAAI/bge-reranker-base
  # link: https://huggingface.co/BAAI/bge-reranker-base
  rerank = SentenceTransformerRerank(
      top_n=2, model="BAAI/bge-reranker-base"
  )
  from llama_index import QueryBundle
  from llama_index.schema import TextNode, NodeWithScore
  query = QueryBundle("I want a dog.")
  scored_nodes = [
      NodeWithScore(node=TextNode(text="This is a cat"), score=0.6),
      NodeWithScore(node=TextNode(text="This is a dog"), score=0.4),
  ]
  reranked_nodes = rerank.postprocess_nodes(
      scored_nodes, query_bundle=query
  )
  ## 6.使用
  

  # 将步骤内容组合起来
  import os
  from llama_index import ServiceContext, VectorStoreIndex, StorageContext
  from llama_index.node_parser import SentenceWindowNodeParser
  from llama_index.indices.postprocessor import MetadataReplacementPostProcessor
  from llama_index.indices.postprocessor import SentenceTransformerRerank
  from llama_index import load_index_from_storage
  
  def build_sentence_window_index(
      documents,
      llm,
      embed_model="local:BAAI/bge-small-en-v1.5",
      sentence_window_size=3,
      save_dir="sentence_index",
  ):
      # create the sentence window node parser w/ default settings
      node_parser = SentenceWindowNodeParser.from_defaults(
          window_size=sentence_window_size,
          window_metadata_key="window",
          original_text_metadata_key="original_text",
      )
      ## 上下文对象，利用它来增强对应的句子块
      sentence_context = ServiceContext.from_defaults(
          llm=llm,
          embed_model=embed_model,
          node_parser=node_parser,
      )
      if not os.path.exists(save_dir):
          sentence_index = VectorStoreIndex.from_documents(
              documents, service_context=sentence_context
          )
          sentence_index.storage_context.persist(persist_dir=save_dir)
      else:
          sentence_index = load_index_from_storage(
              StorageContext.from_defaults(persist_dir=save_dir),
              service_context=sentence_context,
          )
  
      return sentence_index
  
  ## 输入句子块查询
  def get_sentence_window_query_engine(
      sentence_index, similarity_top_k=6, rerank_top_n=2
  ):
      # define postprocessors
      postproc = MetadataReplacementPostProcessor(target_metadata_key="window")
      rerank = SentenceTransformerRerank(
          top_n=rerank_top_n, model="BAAI/bge-reranker-base"
      )
  
      sentence_window_engine = sentence_index.as_query_engine(
          similarity_top_k=similarity_top_k, node_postprocessors=[postproc, rerank]
      )
      return sentence_window_engine
  

  ## 6.使用
  ## 设置相关参数
  from llama_index.llms import OpenAI
  index = build_sentence_window_index(
      [document],
      llm=OpenAI(model="gpt-3.5-turbo", temperature=0.1),
      save_dir="./sentence_index",
  )
  ## 输入句子Query
  query_engine = get_sentence_window_query_engine(index, similarity_top_k=6)
  ## 使用句窗检索
  

• 对句窗检索进行评估

  ## 加载进行评估的问题
  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)
  ## 进行标准问答评估
  from trulens_eval import Tru
  def run_evals(eval_questions, tru_recorder, query_engine):
      for question in eval_questions:
          with tru_recorder as recording:
              response = query_engine.query(question)
  ## 调整句子窗口大小参数
  from utils import get_prebuilt_trulens_recorder
  from trulens_eval import Tru
  Tru().reset_database()
  ### 句窗大小设置为1，并运行
  sentence_index_1 = build_sentence_window_index(
      documents,
      llm=OpenAI(model="gpt-3.5-turbo", temperature=0.1),
      embed_model="local:BAAI/bge-small-en-v1.5",
      sentence_window_size=1,
      save_dir="sentence_index_1",
  )
  sentence_window_engine_1 = get_sentence_window_query_engine(
      sentence_index_1
  )
  tru_recorder_1 = get_prebuilt_trulens_recorder(
      sentence_window_engine_1,
      app_id='sentence window engine 1'
  )
  run_evals(eval_questions, tru_recorder_1, sentence_window_engine_1)
  ### 句窗大小设置为3
  sentence_index_3 = build_sentence_window_index(
      documents,
      llm=OpenAI(model="gpt-3.5-turbo", temperature=0.1),
      embed_model="local:BAAI/bge-small-en-v1.5",
      sentence_window_size=3,
      save_dir="sentence_index_3",
  )
  sentence_window_engine_3 = get_sentence_window_query_engine(
      sentence_index_3
  )
  
  tru_recorder_3 = get_prebuilt_trulens_recorder(
      sentence_window_engine_3,
      app_id='sentence window engine 3'
  )
  ## 运行查看结果
  run_evals(eval_questions, tru_recorder_1, sentence_window_engine_1)