Embedding 与 Reranker 深度调优：检索质量的决定性引擎

Bi-Encoder（召回层，两段式）
  查询: [CLS] query text → 向量 q
  文档: [CLS] doc text  → 向量 d
  相似度 = cosine(q, d)
  特点：文档可离线预向量化，检索快（百万级/百毫秒）
  缺点：查询与文档无交叉交互，精度上限低

Cross-Encoder（精排层，交互式）
  输入: [CLS] query [SEP] doc → Transformer 全交互
  输出: 0~1 相关度分数
  特点：词级交叉注意力，精度高
  缺点：文档不能预计算，逐条打分慢（每秒仅数十~数百条）

查询
  │
  ▼
① Bi-Encoder 召回：向量检索 Top-100（毫秒级）
  │
  ▼
② Query 改写 / 混合检索融合（可选，提高召回面）
  │
  ▼
③ Cross-Encoder 精排：Top-100 → Top-10（秒级）
  │
  ▼
④ 送入 LLM 的最终上下文

# chunking_optimizer.py — 按语义边界分块
def semantic_chunk(text: str, max_tokens: int = 400,
                   overlap: int = 80) -> list[str]:
    """按段落 + 语义边界分块，块间有重叠避免信息断裂。"""
    import re
    paragraphs = [p for p in re.split(r"\n\s*\n", text) if p.strip()]

    chunks, current = [], ""
    for para in paragraphs:
        if token_len(current + para) > max_tokens and current:
            chunks.append(current.strip())
            current = para[-overlap:] + para   # 保留尾部重叠
        else:
            current += "\n" + para
    if current.strip():
        chunks.append(current.strip())
    return chunks

# 评估不同粒度对检索质量的影响
def evaluate_chunk_strategy(docs, queries, embed_fn, reranker,
                            judge_fn):
    strategies = {
        "whole_doc": [d for d in docs],
        "semantic_400": flat([semantic_chunk(d) for d in docs]),
        "sentence": flat([split_sentences(d) for d in docs]),
    }
    results = {}
    for name, units in strategies.items():
        recall = measure_recall(units, queries, embed_fn, reranker)
        results[name] = recall
    return results

# hybrid_search.py — BM25 + 向量混合
def hybrid_search(query, bm25_index, vector_store, embed_fn,
                  weights=(0.4, 0.6), top_k=100) -> list[dict]:
    """
    weights: (bm25 权重, 向量权重)
    归一化各自分数后加权融合。
    """
    # 1. BM25 稀疏检索
    bm25_hits = bm25_index.search(query, k=top_k)
    bm25_scores = {doc_id: norm(s) for doc_id, s in bm25_hits}

    # 2. 向量密集检索
    vec = embed_fn(query)
    vec_hits = vector_store.search(vec, k=top_k)
    vec_scores = {doc_id: norm(s) for doc_id, s in vec_hits}

    # 3. 融合排序（RRF 或加权和）
    merged_ids = set(bm25_scores) | set(vec_scores)
    fused = []
    for doc_id in merged_ids:
        s = weights[0] * bm25_scores.get(doc_id, 0) + \
            weights[1] * vec_scores.get(doc_id, 0)
        fused.append({"doc_id": doc_id, "score": s})
    fused.sort(key=lambda x: x["score"], reverse=True)
    return fused[:top_k]

def rrf_fusion(ranked_lists: list[list[str]], k: int = 60) -> list[str]:
    """多路排序列表的 RRF 融合：1/(k + rank) 累加。"""
    from collections import defaultdict
    score = defaultdict(float)
    for ranked in ranked_lists:
        for rank, doc_id in enumerate(ranked):
            score[doc_id] += 1.0 / (k + rank + 1)
    return sorted(score, key=score.get, reverse=True)

def calibrate_hybrid_weights(bm25_index, vector_store, embed_fn,
                             queries, relevant_sets,
                             candidates=[0.3, 0.5, 0.7]) -> float:
    """用标注数据搜索最优 BM25 权重。"""
    best_w, best_score = None, 0
    for w in candidates:
        avg = mean([
            hits_at_k(hybrid_search(q, bm25_index, vector_store, embed_fn,
                                    weights=(w, 1 - w)),
                      relevant_sets[i], k=10)
            for i, q in enumerate(queries)])
        if avg > best_score:
            best_w, best_score = w, avg
    return best_w

# query_rewriter.py — 查询改写
REWRITE_PROMPT = """将用户查询改写为更适合文档检索的形式：
- 补全指代（"它"→具体对象）
- 使用书面、具体的检索词
- 保留原始意图，不改变语义
- 若已是清晰检索式查询则原样返回
用户查询：{query}
改写结果（只输出改写文本）："""

def rewrite_query(query, history, llm_call) -> str:
    # 多轮对话中，把"它"/"这个"等指代结合上文解析
    if history:
        contextual = f"上文：{history[-1]}\n当前：{query}"
        rewritten = llm_call(REWRITE_PROMPT.format(query=contextual))
    else:
        rewritten = llm_call(REWRITE_PROMPT.format(query=query))
    return rewritten.strip() or query   # 空结果回退原查询

def hyde_retrieve(query, embed_fn, vector_store, llm_call, top_k=20):
    """Hypothetical Document Embeddings：先假答后检索。"""
    hypo_doc = llm_call(
        f"写一段包含可能答案的科普性短文，用于检索参考：{query}")
    hypo_vec = embed_fn(hypo_doc)
    # 用假设文档的向量检索真实文档
    return vector_store.search(hypo_vec, k=top_k)

def expand_query(query, llm_call) -> list[str]:
    """生成同义变体查询，多路召回后合并。"""
    variants = llm_call(
        f"为查询生成 3 个同义变体，只输出变体，每行一个：{query}")
    return [query] + [v.strip() for v in variants.splitlines() if v.strip()]

def multi_query_recall(query, embed_fn, vector_store, llm_call,
                       top_k_per=20) -> list[str]:
    """多查询召回：各路结果 RRF 合并。"""
    all_ranked = []
    for variant in expand_query(query, llm_call):
        vec = embed_fn(variant)
        hits = vector_store.search(vec, k=top_k_per)
        all_ranked.append([h["id"] for h in hits])
    return rrf_fusion(all_ranked)

# reranker.py — 交叉编码器精排
from sentence_transformers import CrossEncoder

class Reranker:
    def __init__(self, model_name="BAAI/bge-reranker-v2-m3"):
        self.model = CrossEncoder(model_name)

    def rerank(self, query: str, candidates: list[dict], top_k: int = 10
               ) -> list[dict]:
        """对候选文档逐一打分并重排。"""
        pairs = [(query, c["text"]) for c in candidates]
        scores = self.model.predict(pairs, batch_size=8)
        # 返回带分数的重排结果
        ranked = sorted(zip(candidates, scores),
                        key=lambda x: x[1], reverse=True)
        return [{**c, "rerank_score": float(s)}
                for c, s in ranked[:top_k]]


# 完整流水线：召回 → 精排
def retrieve_and_rerank(query, vector_store, bm25, embed_fn,
                        reranker, top_recall=100, top_final=10):
    recalled = hybrid_search(query, bm25, vector_store, embed_fn,
                             top_k=top_recall)
    final = reranker.rerank(query, recalled, top_k=top_final)
    return final

def measure_rerank_gain(queries, vector_store, reranker,
                        relevant_sets):
    """对比精排前后的 NDCG 提升。"""
    gains = []
    for query, relevant in zip(queries, relevant_sets):
        recalled = vector_store.search(embed_fn(query), k=100)
        ndcg_before = ndcg_at_k(recalled, relevant, k=10)
        reranked = reranker.rerank(query, recalled, top_k=10)
        ndcg_after = ndcg_at_k(reranked, relevant, k=10)
        gains.append(ndcg_after - ndcg_before)
    avg_gain = mean(gains)
    print(f"NDCG@10 平均提升: {avg_gain:+.3f}")
    return avg_gain

# train_data_builder.py — 构造训练数据
def build_training_pairs(corpus, queries, positive_negative_fn) -> list[dict]:
    """构造 (query, doc, label) 三元组，label∈{0,1}。"""
    pairs = []
    for query in queries:
        pos = positive_negative_fn(query)["positive"]    # 相关文档
        neg = positive_negative_fn(query)["negative"]    # 难负样本
        pairs.append({"query": query, "doc": pos, "label": 1})
        for n in neg[:3]:
            pairs.append({"query": query, "doc": n, "label": 0})
    return pairs

def mine_hard_negatives(query, vector_store, embed_fn,
                        relevant_ids: set, top_k=50) -> list[str]:
    """向量检索得分高但并非相关的文档 = 难负样本。"""
    hits = vector_store.search(embed_fn(query), k=top_k)
    return [h["id"] for h in hits if h["id"] not in relevant_ids][:20]

# finetune_reranker.py — 微调 Cross-Encoder
from sentence_transformers import CrossEncoder, losses
from sentence_transformers.models import Transformer
import torch

def finetune_reranker(base_model="BAAI/bge-reranker-v2-m3",
                      train_pairs: list[dict],
                      output_dir="models/reranker-domain"):
    """在领域数据上微调 reranker。"""
    model = CrossEncoder(base_model, num_labels=1)

    samples = [(p["query"], p["doc"], float(p["label"]))
               for p in train_pairs]
    model.fit(
        train_dataloader=samples,
        loss=losses.MultipleNegativesRankingLoss(model),
        epochs=3,
        warmup_steps=100,
        output_path=output_dir,
    )
    print(f"已微调完成，保存至 {output_dir}")

def validate_finetuned(base_reranker, finetuned, eval_queries,
                       relevant_sets):
    """对比基础与微调 reranker 在留存集上的 NDCG。"""
    base_ndcg = mean(ndcg_at_k(base_reranker.rerank(q, corpus, 10), rel, 10)
                     for q, rel in zip(eval_queries, relevant_sets))
    ft_ndcg = mean(ndcg_at_k(finetuned.rerank(q, corpus, 10), rel, 10)
                   for q, rel in zip(eval_queries, relevant_sets))
    print(f"基础 {base_ndcg:.3f} → 微调 {ft_ndcg:.3f}")
    assert ft_ndcg >= base_ndcg, "微调没有提升，检查训练数据质量"
    return ft_ndcg - base_ndcg

def build_ir_eval_set(corpus, queries, annotator) -> list[dict]:
    """构造检索评测集：每个查询标注相关文档集合 + 分级。"""
    eval_cases = []
    for query in queries:
        relevant = annotator(query, corpus)   # 人工或 LLM 标注
        eval_cases.append({
            "query": query,
            "relevant_ids": {r["id"] for r in relevant},
            "grades": {r["id"]: r["grade"] for r in relevant},  # 0/1/2 分级
        })
    return eval_cases

def run_ir_eval(retrieval_system, eval_cases, k_list=[5, 10, 20]):
    """输出完整检索评测报告。"""
    report = {}
    for k in k_list:
        recalls = [recall_at_k(retrieval_system(q), rel, k)
                   for q, rel in eval_cases]
        report[f"recall@{k}"] = round(mean(recalls), 3)
    ndcgs = [ndcg_at_k(retrieval_system(q), grades, 10)
             for q, grades in eval_cases]
    report["ndcg@10"] = round(mean(ndcgs), 3)
    return report

def end_to_end_compare(retrievers: dict, golden, judge_fn):
    """对比不同检索方案对最终 LLM 回答质量的影响。"""
    results = {}
    for name, retriever in retrievers.items():
        scores = []
        for case in golden.cases:
            evidence = retriever(case["query"])
            answer = llm_with_context(case["query"], evidence)
            scores.append(judge_fn(case["query"], answer,
                                   case["reference"])["total"])
        results[name] = mean(scores)
    return results   # 例：{"vector_only": 17.2, "hybrid+rerank": 20.1}

# 索引参数影响召回与延迟，需做基准测试
def benchmark_index(vector_store, queries, relevant, index_configs):
    for cfg in index_configs:
        vector_store.configure_index(cfg)   # {"type":"hnsw","M":16,"efConstruction":100}
        lat = measure_latency(vector_store, queries)
        rec = recall_at_k(vector_store, queries, relevant, k=10)
        print(f"{cfg}: 延迟 {lat:.0f}ms, 召回 {rec:.3f}")

def profile_pipeline(query, vector_store, reranker):
    import time
    t0 = time.time(); recalled = vector_store.search(query_vec, k=100)
    t1 = time.time(); final = reranker.rerank(query, recalled, top_k=10)
    t2 = time.time()
    print(f"召回 {1000*(t1-t0):.0f}ms | 精排 {1000*(t2-t1):.0f}ms")
    # 若精排占比 > 70%，优先缩小候选量或用更小模型

def monitor_retrieval(query, final_evidence, click_or_feedback):
    """记录检索命中与用户反馈，累积检索质量指标。"""
    push_metric("retrieval.top1_score", final_evidence[0]["score"])
    push_metric("retrieval.count", len(final_evidence))
    if click_or_feedback is not None:
        push_metric("retrieval.satisfied", click_or_feedback)

def multilingual_retrieve(query, vec_store, embed_fn, reranker, top_k=10):
    # 查询可能是中文，文档含英文——多语言模型统一嵌入空间
    vec = embed_fn(query)   # 如 bge-m3 / multilingual-e5
    candidates = vec_store.search(vec, k=100)
    # 跨语言精排仍有效（Cross-Encoder 跨语言交互）
    return reranker.rerank(query, candidates, top_k)

def graph_aware_retrieve(query, entity_extractor, graph_search,
                         vector_store, embed_fn, reranker):
    entities = entity_extractor(query)
    graph_hits = graph_search(entities[0]) if entities else []
    vec_hits = vector_store.search(embed_fn(query), k=50)
    merged = rrf_fusion([graph_hits, [h["id"] for h in vec_hits]])
    return reranker.rerank(query, load_docs(merged), top_k=10)

def llm_annotate_relevance(query, doc, judge_llm) -> int:
    """LLM 判断文档与查询的相关度（0/1/2），用于扩大评测集。"""
    verdict = judge_llm(
        f"文档与查询的相关度？0=无关,1=部分,2=高度相关。"
        f"\n查询：{query}\n文档：{doc[:500]}",
        output_schema={"type": "integer", "enum": [0, 1, 2]})
    return verdict

# 人工抽检 10% 保证标注质量
def human_audit(annotated_cases, audit_ratio=0.1):
    sample = random.sample(annotated_cases, int(len(annotated_cases) * audit_ratio))
    agreements = [annotate_again_manually(c) == c["label"] for c in sample]
    return mean(agreements)   # 一致性 > 0.9 认为标注可信