RAG系统三阶优化实录：分块、Embedding与Rerank的量化调优

from langchain.text_splitter import TextSplitter
from typing import List, Dict
import re

class StructureAwareSplitter(TextSplitter):
    def __init__(self, chunk_size: int = 700, chunk_overlap: int = 100):
        super().__init__(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
        self.table_pattern = re.compile(r'^\|.*\|$', re.MULTILINE)

    def split_text(self, text: str) -> List[str]:
        # 第一步：按三级标题切分
        sections = re.split(r'(?m)^(?=####?\s)', text)
        chunks = []
        for section in sections:
            if not section.strip():
                continue
            # 第二步：处理段内表格——若表格行数>3，整表捕捉；否则按行合并到邻近文本
            if self.table_pattern.search(section):
                lines = section.split('\n')
                buffer = []
                for line in lines:
                    buffer.append(line)
                    if self.table_pattern.match(line) and len(buffer) >= 5:
                        # 遇到连续表格行，则flush成一个chunk
                        merged = '\n'.join(buffer)
                        if len(merged)  self.chunk_size:
                    parts = re.split(r'(? self.chunk_size:
                            chunks.append(temp)
                            temp = part
                        else:
                            temp += part
                    if temp:
                        chunks.append(temp)
                else:
                    chunks.append(section)
        return chunks

# 实例化并覆盖原splitter
splitter = StructureAwareSplitter(chunk_size=700, chunk_overlap=100)

from FlagEmbedding import FlagModel
import numpy as np
from langchain.embeddings.base import Embeddings

class GTEQwen2Embeddings(Embeddings):
    def __init__(self, model_path: str = "/models/GTE_Qwen2-7B-instruct"):
        self.model = FlagModel(
            model_path,
            query_instruction_for_retrieval="为检索到相关段落，请查询：",
            use_fp16=True,  # 7B模型必须用fp16，否则OOM
            devices="cuda:0"
        )

    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        # 文档端不需要指令前缀
        return self.model.encode(texts, batch_size=8).tolist()

    def embed_query(self, text: str) -> List[float]:
        # query端自动附加指令
        return self.model.encode_queries([text], batch_size=1).tolist()[0]

# 注意：langchain的HuggingFaceEmbeddings无法直接加载FlagModel，必须自定义类

from FlagEmbedding import FlagReranker

reranker = FlagReranker('/models/bge-reranker-v2-m3', use_fp16=True)

def rerank(query: str, docs: List[str], top_k: int = 3):
    # 构造pair列表
    pairs = [(query, doc) for doc in docs]
    scores = reranker.compute_score(pairs, normalize=True)  # normalize=True返回0~1概率
    # 按分数降序取top_k
    sorted_idx = np.argsort(scores)[::-1][:top_k]
    return [docs[i] for i in sorted_idx], [scores[i] for i in sorted_idx]

# 在检索链中集成：
# 1. 用GTE召回top_20 = vector_store.similarity_search(query, k=20)
# 2. final_docs = rerank(query, [d.page_content for d in top_20], top_k=3)