RAG系统优化实录：chunk调整、embedding切换与rerank引入

# chunk_strategy.py
import re
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.schema import Document

def semantic_chunk(documents: list[Document], chunk_size=800, chunk_overlap=100):
    """
    语义分割：先按Markdown标题分割，再按段落分割，最后用RecursiveCharacterTextSplitter分块
    """
    # 第一阶段：按Markdown标题分割
    heading_splitter = RecursiveCharacterTextSplitter(
        separators=["\n## ", "\n### ", "\n#### ", "\n\n", "\n", " "],
        chunk_size=2000,  # 先设置一个较大的值，只做标题级分割
        chunk_overlap=0,
        length_function=len,
    )

    heading_docs = heading_splitter.split_documents(documents)

    # 第二阶段：对每个标题块再做细粒度分块
    final_splitter = RecursiveCharacterTextSplitter(
        separators=["\n\n", "\n", "。", ".", " ", ""],
        chunk_size=chunk_size,
        chunk_overlap=chunk_overlap,
        length_function=len,
        keep_separator=True,
    )

    final_docs = final_splitter.split_documents(heading_docs)

    # 后处理：合并小chunk
    merged_docs = []
    buffer = ""
    for doc in final_docs:
        if len(buffer + doc.page_content) < 200:
            buffer += doc.page_content
        else:
            if buffer:
                merged_docs.append(Document(page_content=buffer, metadata=doc.metadata))
            buffer = doc.page_content
    if buffer:
        merged_docs.append(Document(page_content=buffer, metadata=final_docs[-1].metadata))

    return merged_docs

# 使用示例
with open("network_docs.md", "r", encoding="utf-8") as f:
    raw_text = f.read()
raw_docs = [Document(page_content=raw_text, metadata={"source": "network_docs.md"})]
chunks = semantic_chunk(raw_docs, chunk_size=800, chunk_overlap=100)
print(f"原始文档大小: {len(raw_text)}字符，生成chunk数: {len(chunks)}")

# rag_optimizer.py
import os
from langchain_community.embeddings import OpenAIEmbeddings
from langchain_community.embeddings import HuggingFaceBgeEmbeddings
from langchain_community.vectorstores import Chroma
from langchain.retrievers import ContextualCompressionRetriever
from langchain.retrievers.document_compressors import CrossEncoderReranker
from langchain_community.cross_encoders import HuggingFaceCrossEncoder

# 配置参数
EMBEDDING_MODEL = "BAAI/bge-large-zh-v1.5"  # 切换为国产模型
RERANK_MODEL = "BAAI/bge-reranker-v2-m3"
VECTOR_DB_PATH = "./chroma_db_optimized"
CHUNK_SIZE = 800
CHUNK_OVERLAP = 100

# 初始化embedding模型
embedding_model = HuggingFaceBgeEmbeddings(
    model_name=EMBEDDING_MODEL,
    model_kwargs={"device": "cuda"},  # 如果有GPU，使用CUDA加速
    encode_kwargs={
        "normalize_embeddings": True,  # 归一化后余弦相似度等价于内积
        "batch_size": 32,
    },
    query_instruction="为这个句子生成表示以用于检索相关文章：",  # bge模型的prompt模板
)

# 构建向量库
vectorstore = Chroma.from_documents(
    documents=chunks,
    embedding=embedding_model,
    persist_directory=VECTOR_DB_PATH,
    collection_metadata={"hnsw:space": "cosine"},  # 使用余弦距离
)

# 基础检索器（仅向量检索）
base_retriever = vectorstore.as_retriever(
    search_type="similarity",
    search_kwargs={"k": 20}  # 先取Top-20给rerank筛选
)

# 初始化rerank模型
cross_encoder = HuggingFaceCrossEncoder(model_name=RERANK_MODEL)
compressor = CrossEncoderReranker(
    model=cross_encoder,
    top_n=5,  # 最终返回Top-5
    batch_size=16,
)

# 组合检索器：向量检索 + rerank重排序
optimized_retriever = ContextualCompressionRetriever(
    base_compressor=compressor,
    base_retriever=base_retriever,
)

# 测试查询
query = "OSPF多区域配置步骤"
docs = optimized_retriever.get_relevant_documents(query)
for i, doc in enumerate(docs):
    print(f"Rank {i+1}: {doc.page_content[:100]}... (score: {doc.metadata.get('relevance_score', 'N/A')})")