RAG系统优化：chunk策略、Embedding切换与Rerank引入实践

# chunk_optimizer.py
from langchain.text_splitter import RecursiveCharacterTextSplitter

def build_splitter(chunk_size: int = 512, chunk_overlap: int = 80):
    """
    使用RecursiveCharacterTextSplitter，按段落、句子、单词层级切分
    适合中文技术文档：优先保留段落完整性
    """
    separators = ["\n\n", "\n", "。", ".", " ", ""]
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=chunk_size,
        chunk_overlap=chunk_overlap,
        separators=separators,
        length_function=len,  # 使用字符数而非token数，更保守
        is_separator_regex=False
    )
    return text_splitter

# 使用示例
splitter = build_splitter(chunk_size=512, chunk_overlap=80)
docs = splitter.split_documents(raw_documents)
print(f"原始文档数: {len(raw_documents)} -> chunk数: {len(docs)}")

# embedding_switch.py
from sentence_transformers import SentenceTransformer
from langchain.embeddings import HuggingFaceBgeEmbeddings

model_name = "BAAI/bge-large-zh-v1.5"
model_kwargs = {'device': 'cuda:0'}
encode_kwargs = {'normalize_embeddings': True}  # 余弦相似度需要归一化

bge_embeddings = HuggingFaceBgeEmbeddings(
    model_name=model_name,
    model_kwargs=model_kwargs,
    encode_kwargs=encode_kwargs,
    query_instruction="为这个句子生成表示以用于检索相关文章：",  # bge官方前缀
    embed_instruction="",
)

# 注意：LangChain的HuggingFaceBgeEmbeddings会自动处理前缀
# 但如果你直接调用model.encode()，需要手动加前缀
# 测试相似度：
query = "微服务熔断降级机制"
passage = "当服务调用失败率达到阈值时，熔断器打开，后续请求直接返回降级结果"
query_vec = bge_embeddings.embed_query(query)
passage_vec = bge_embeddings.embed_documents([passage])[0]
from sklearn.metrics.pairwise import cosine_similarity
import numpy as np
sim = cosine_similarity([query_vec], [passage_vec])[0][0]
print(f"bge-large 相似度: {sim:.4f}")  # 输出0.9234
# 而之前ada-002只有0.72

# reranker.py
from transformers import AutoModelForSequenceClassification, AutoTokenizer
import torch

class BGEReranker:
    def __init__(self, model_name: str = "BAAI/bge-reranker-v2-m3"):
        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.model = AutoModelForSequenceClassification.from_pretrained(
            model_name,
            trust_remote_code=True,
            torch_dtype=torch.float16  # 减少显存占用
        ).to(self.device)
        self.model.eval()

    def rerank(self, query: str, passages: list[str], top_k: int = 5) -> list[tuple[int, float]]:
        pairs = [[query, passage] for passage in passages]
        with torch.no_grad():
            inputs = self.tokenizer(
                pairs,
                padding=True,
                truncation=True,
                max_length=512,
                return_tensors="pt"
            ).to(self.device)
            scores = self.model(**inputs).logits.squeeze(-1).cpu().numpy()
        # 按分数降序排列
        sorted_indices = scores.argsort()[::-1][:top_k]
        return [(idx, scores[idx]) for idx in sorted_indices]

# 使用示例
reranker = BGEReranker()
query = "如何配置Hystrix线程池大小？"
candidates = chroma_collection.similarity_search(query, k=30)  # 先召回30个
passages = [doc.page_content for doc in candidates]
reranked = reranker.rerank(query, passages, top_k=5)
for idx, score in reranked:
    print(f"分数: {score:.4f}, 内容: {passages[idx][:60]}...")