向量化与重排模型服务优化：Embedding 与 Reranker 的性能工程

关键差异：
Embedding：N 段文本 → N 个向量（可并行、可缓存、可离线批量）
Reranker：query × M 个候选 → M 个分数（在线、算力密集）

→ Embedding 优化 = 吞吐工程（批大、量化、多流）
→ Reranker 优化 = 延迟工程（裁剪候选、批内并行、蒸馏）

# 批处理编码（GPU 利用率随 batch 提升）
import torch
from transformers import AutoTokenizer, AutoModel

tok = AutoTokenizer.from_pretrained("BAAI/bge-large-zh-v1.5")
model = AutoModel.from_pretrained("BAAI/bge-large-zh-v1.5").cuda().half()

@torch.inference_mode()
def encode(texts, batch_size=256, max_len=512):
    out = []
    for i in range(0, len(texts), batch_size):
        batch = texts[i:i+batch_size]
        enc = tok(batch, padding=True, truncation=True,
                  max_length=max_len, return_tensors="pt").to("cuda")
        hidden = model(**enc).last_hidden_state
        # BGE 用 CLS token，并做 L2 归一化
        vec = torch.nn.functional.normalize(hidden[:, 0], dim=-1)
        out.append(vec.cpu())
    return torch.cat(out)

批处理调优要点：
□ batch_size 越大吞吐越高，直到显存或延迟拐点
□ 按长度分桶（length bucketing）→ 减少 padding 浪费
□ 排序后分批（相似长度聚在一起）→ 提升有效算力
□ padding 到批内最长即可，别固定到 max_len

量化收益：
□ FP16：显存减半，编码提速 ~1.5-2x，精度几乎无损
□ INT8（动态量化）：再提速，精度损失小
□ 对 Embedding 这类「输出是连续向量」的模型，
  INT8 量化对检索质量影响通常可接受（需评测验证）

注意：
□ 量化后必须重新评测检索指标（Recall@K、MRR）
□ 不同模型对量化的敏感度差异大，别一刀切

# CPU 上用 ONNX Runtime + INT8 动态量化
python -m onnxruntime.quantization.preprocess \
    --input model.onnx --output model_pre.onnx
python -m onnxruntime.quantization.quantize_dynamic \
    --input model_pre.onnx --output model_int8.onnx \
    --weight_type QInt8

成本模型：
Reranker 延迟 ≈ 候选数 M × 单对计算时间
单对计算时间 ∝ 序列长度（query + 文档）

→ 降延迟的两条路：
① 减少 M（候选裁剪）：先用便宜手段把 M 从 100 降到 20
② 降低单对成本（量化、蒸馏、短序列）

两阶段检索：
Stage 1  向量检索（Embedding）→ Top-100 候选（便宜、快）
Stage 2  Reranker 精排 → Top-10（贵、准）
Stage 3  送入 LLM 生成

好处：
□ Reranker 只需处理 100 个候选，而非全库
□ 用 Embedding 的召回能力 + Reranker 的精度
□ 延迟可控（精排规模固定）

# 两阶段检索示意
candidates = vector_search(query_emb, top_k=100)      # 粗排
pairs = [(query, c.text) for c in candidates]
scores = reranker.predict(pairs, batch_size=32)       # 精排
top = [c for c, s in sorted(zip(candidates, scores),
                            key=lambda x: -x[1])[:10]]

Reranker 批处理要点：
□ batch_size 受「对长度 × 批大小」的显存限制
□ 按 query+doc 总长度分桶
□ 长文档先截断（通常取前 512 token 足够）
□ FP16 加速；INT8 需评测精度
□ 蒸馏：用大 Reranker 蒸馏小模型（如 6 层），延迟减半

# Reranker 批量打分（长度分桶）
def rerank(query, docs, batch_size=32, max_len=512):
    pairs = [(query, d) for d in docs]
    scores = []
    for i in range(0, len(pairs), batch_size):
        batch = pairs[i:i+batch_size]
        enc = tok(batch, padding=True, truncation=True,
                  max_length=max_len, return_tensors="pt").to("cuda")
        with torch.inference_mode():
            logits = model(**enc).logits.view(-1)
        scores.extend(logits.float().cpu().tolist())
    return scores

两种部署形态：
□ 独立服务：Embedding 与 Reranker 各自成服务（推荐）
  → 独立扩缩容、独立选型、故障隔离
□ 合体服务：一个进程同时提供两者
  → 省资源，但耦合（一个慢拖累另一个）

接口设计：
POST /embed   {"texts": [...]} → {"vectors": [...]}
POST /rerank  {"query": "...", "docs": [...]} → {"scores": [...]}
□ 都支持批处理（一次传多段/多对）
□ 都返回 usage（用于计费与监控）

Embedding 缓存：
□ 文档向量：入库时算一次，永久缓存（内容哈希为 key）
□ query 向量：相同 query 命中缓存（精确 + 语义）
□ 命中率通常很高（重复查询、重复文档）

Reranker 缓存：
□ (query, doc) 对分数：可缓存，但组合爆炸 → 收益有限
□ 更适合缓存「最终 Top-K 结果」而非中间分数

# 缓存 key 设计：内容哈希（而非文档 ID）
import hashlib
def cache_key(text, model_name):
    h = hashlib.sha256(text.encode()).hexdigest()[:16]
    return f"emb:{model_name}:{h}"     # 换模型即失效，避免串味

# 命中则跳过编码；未命中则批量编码后回填
def embed_cached(texts, store, model_name):
    keys = [cache_key(t, model_name) for t in texts]
    hits = store.mget(keys)
    miss_idx = [i for i, v in enumerate(hits) if v is None]
    if miss_idx:
        vecs = encode([texts[i] for i in miss_idx])
        for i, v in zip(miss_idx, vecs):
            store.set(keys[i], v, ex=30*24*3600)   # 30 天过期
            hits[i] = v
    return hits

检索栈与生成栈的协同：
□ 检索服务独立扩缩容（检索 QPS 与生成 QPS 不同）
□ 检索延迟纳入端到端 SLO（TTFT 的一部分）
□ 检索服务故障时的降级（返回缓存/跳过精排）
□ 共享 GPU 池时做优先级隔离（见 GPU 共享与调度）

检索质量指标：
□ Recall@K：Top-K 里包含相关文档的比例（召回）
□ MRR：首个相关文档的排名倒数均值
□ NDCG@K：考虑排序位置的相关性
□ 端到端：RAG 答案的忠实度/正确率（最终标准）

性能指标：
□ Embedding：texts/s（吞吐）、P99 单批延迟
□ Reranker：pairs/s、P99 精排延迟
□ 缓存命中率

评测纪律：
□ 量化/蒸馏前后必须跑同一评测集对比
□ 关注「质量下降是否可接受」，而非「是否有下降」
□ 质量-延迟曲线：找 SLO 内的最优工作点