告别碎片化！两大先进分块技术如何提升RAG的语义连贯性？

完整文档 → Token级嵌入 → 分块 → 池化 → 块嵌入

分块 → LLM生成上下文 → 块+上下文嵌入 → 混合检索（BM25+密集检索） → 重排序

"基于以下文档，为当前块生成上下文摘要：[文档内容] [当前块]"

class BgeEmbedder():
  def __init__(self):
    self.model_id = 'BAAI/bge-m3'
    self.model = AutoModel.from_pretrained(self.model_id, trust_remote_code=True).cuda();
    self.tokenizer = AutoTokenizer.from_pretrained(self.model_id);
    self.model.eval()



  def get_model(self):
    return self.model

  def encode(self, chunks, task, max_length=512):
    if type( chunks) == str:
      chunks = [chunks]
    chunks_embeddings = []
    for chunk in chunks:
      tokens = self.tokenizer(chunk, return_tensors='pt')['input_ids'].to("cuda")
      with torch.no_grad():
        token_embeddings = self.model(input_ids=tokens)[0]
        pooled_embeddings = token_embeddings.sum(dim=1) / len(token_embeddings)
        pooled_embeddings = F.normalize(pooled_embeddings, p=2, dim=1).squeeze(0)
      chunks_embeddings.append(pooled_embeddings.to(dtype=torch.float32).detach().cpu().numpy())
    if task == 'retrieval.query':
      return chunks_embeddings
    else:
      return [chunks_embeddings]

  # def encode(self, chunks, task, max_length=512):
  #   if type( chunks) == str:
  #     chunks = [chunks]
  #   chunks_embeddings = []
  #   for chunk in chunks:
  #     tokens = self.tokenizer(chunk, return_tensors='pt')['input_ids'].to("cuda")
  #     with torch.no_grad():
  #       token_embeddings = self.model(input_ids=tokens)[0]
  #       div = token_embeddings.size(1)
  #       pooled_embeddings = token_embeddings.sum(dim=1) / div
  #       pooled_embeddings = F.normalize(pooled_embeddings, p=2, dim=1).squeeze(0)
  #     chunks_embeddings.append(pooled_embeddings.to(dtype=torch.float32).detach().cpu().numpy())
  #   if task == 'retrieval.query':
  #     return chunks_embeddings
  #   else:
  #     return [chunks_embeddings]


  def encodeLateChunking(self, text_token, span_annotation, task, max_length = None ):
    with torch.no_grad():
      token_embeddings = self.model(input_ids=text_token)[0]
    outputs = []
    for embeddings, annotations in zip(token_embeddings, span_annotation):
        if (max_length is not None):  # remove annotations which go bejond the max-length of the model
            annotations   = [
                              (start, min(end, max_length - 1))
                                for (start, end) in annotations
                                  if start < (max_length - 1)
                            ]
        pooled_embeddings = [
                              embeddings[start:end].sum(dim=0) / (end - start)
                                for (start, end) in annotations
                                  if (end - start) >= 1
                            ]
        pooled_embeddings = [
                              F.normalize(pooled_embedding.unsqueeze(-1), p=2, dim=0).squeeze(-1)
                                for  pooled_embedding in pooled_embeddings
                            ]
        pooled_embeddings = [
                              embedding.to(dtype=torch.float32).detach().cpu().numpy()
                                for embedding in pooled_embeddings
                            ]
        outputs.append(pooled_embeddings)
    return outputs

class ChunkContextualizer():
  """
  This class is used to contextualize the chunks of text.
  It gives to the model the document and the chunk of text to generate the context to be added to the chunk.
  The model is a language model that generates the context.
  """
  def __init__(self, corpus, prompt, examples, llm_id, segmenter,load_in_4bit = True): #model_template
      self.corpus = corpus
      self.prompt = prompt
      self.examples = examples
      self.segmenter = segmenter
      #self.model_template = model_template
      #self.chat_template = (model_template,"<|end|>")
      self.llm_id = llm_id
      self.load_in_4bit = load_in_4bit


      dtype = None
      self.llm_model, self.llm_tokenizer = FastLanguageModel.from_pretrained(
                                                                model_name = model_id,
                                                                dtype = dtype,
                                                                load_in_4bit = load_in_4bit,
                                                                )
      self.llm_tokenizer = get_chat_template(self.llm_tokenizer,
                                            chat_template = 'phi-3.5',#self.chat_template,            # apply phi-3.5 chat template
                                            mapping = {"role" : "from", "content" : "value", "user" : "human", "assistant" : "gpt"},
                                            )
      FastLanguageModel.for_inference(self.llm_model);

      self.generation_config =  GenerationConfig(eos_token_id=self.llm_tokenizer.convert_tokens_to_ids("<|end|>"))

  def get_llm(self):
    return self.llm_model

  def get_formatted_prompt(self,doc,chunk):
    ex = '\n'.join([f'CHUNK: {chunk}\nCONTEXT: {context}' for (chunk, context) in self.examples]) # format examples with chunk and context
    return self.prompt.format(doc_content=doc, chunk_content=chunk, examples=ex)                  # format prompt with chunk, document and examples

  def extract_llm_responses(self, conversation: str):                                             # extract assistant responses from the conversation
      # Regex to match assistant output between <|assistant|> and <|end|>
      assistant_responses = re.findall(r"<\|assistant\|>(.*?)<\|end\|>", conversation, re.DOTALL)
      return assistant_responses

  def contextualize(self, chunking_size):
    ids,texts = zip(*[ (key, value['text']) for (key,value) in self.corpus.items() ])
    ids,texts = list(ids), list(texts)
    chunks_ids,contextualized_chunks=[],[]
    pbar = tqdm(total=len(texts), desc="Contextualizing Chunks", mininterval=60.0)
    for i in range(len(texts)):
      chunks = self.segmenter.segment(input_text=texts[i],max_chunk_length=chunking_size)         # segment text into chunks with passed segmenter
      for j in range(len(chunks)):
        chunk = chunks[j]
        doc = texts[i]                                                        # getting document for all the chunks
        prompt = self.get_formatted_prompt(doc,chunk)                         # formatting prompt with current document and current chunk
        messages=[{"from": "human", "value": prompt}]

        inputs = self.llm_tokenizer.apply_chat_template(                      # apply chat template to the prompt to get the input
                                                        messages,
                                                        tokenize=True,
                                                        return_tensors="pt",
                                                        padding=True,
                                                        ).to("cuda")

        output_ids = self.llm_model.generate(
                                              input_ids=inputs,
                                              do_sample=False,
                                              max_new_tokens=2000,
                                              generation_config=self.generation_config,
                                              )

        output = self.llm_tokenizer.batch_decode(output_ids)                  # decode the output to text
        generated_context = self.extract_llm_responses(output[0])             # extract context from output string
        if generated_context:                                                 # only add non-empty entries
            contextualized_chunks.append(chunk + '\n' + generated_context[0]) # append generated context to current chunk
            chunks_ids.append(f'{ids[i]}:{j}')
      pbar.update(1)
    return contextualized_chunks, chunks_ids

contextual_retrieval_prompt = """
DOCUMENT:
{doc_content}


Here is the chunk we want to situate within the whole document
CHUNK:
{chunk_content}


Provide a concise context for the following chunk, focusing on its relevance within the overall document for the purposes of improving search retrieval of the chunk.
Answer only with the generated context and nothing else.
Output your answer after the phrase "The document provides an overview".
Think step by step.


Here there is an example you can look at.
EXAMPLE:
{examples}"""