perf: rerank修复 + MMR向量复用优化

1. jina_rerank: 修复 API key 从环境变量读取,修复响应格式
   (document对象→文本字符串),支持 bge-reranker-v2-m3
2. recall: rerank fallback 编码的向量缓存到 candidate._cached_vector,
   MMR 自动复用(避免 fallback+MMR 双重 encode_batch)
3. recall: 移除重复的 doc_texts 定义(修复语法错误)
This commit is contained in:
小唯 A06
2026-05-27 15:38:04 +08:00
parent d424c443f0
commit 83776df7d3
2 changed files with 48 additions and 15 deletions
+36 -11
View File
@@ -235,7 +235,8 @@ def _get_recall_counts(doc_ids: list[str]) -> dict[str, int]:
def _mmr_diversity(candidates: list[dict], query_vector: list[float],
lambda_mult: float = 0.5, top_k: int = None) -> list[dict]:
lambda_mult: float = 0.5, top_k: int = None,
precomputed_vectors: list = None) -> list[dict]:
"""MMR(Maximum Marginal Relevance)多样性重排序。
Args:
@@ -244,6 +245,8 @@ def _mmr_diversity(candidates: list[dict], query_vector: list[float],
lambda_mult: 多样性权重因子,[0,1],越大越倾向多样性
0=纯相关性排序,0.5=平衡,1=最大多样性
top_k: 返回数量,默认 len(candidates)
precomputed_vectors: 可选,预计算好的文档向量(避免重复 BGE 编码)
如果传入 None,则自动从 doc_meta.text 编码
"""
if not candidates or lambda_mult <= 0:
return candidates[:top_k]
@@ -256,14 +259,25 @@ def _mmr_diversity(candidates: list[dict], query_vector: list[float],
if len(candidates) <= top_k:
return candidates
# 预计算所有候选向量
try:
bge = _get_bge()
texts = [c['doc_meta']['text'] for c in candidates]
vectors = bge.encode_batch(texts, normalize_output=True)
except Exception as e:
print(f"[MMR] 向量编码失败: {e}")
return candidates[:top_k]
# ══ 优化4:复用预计算向量(避免 MMR 重复 encode)═
# 支持从 rerank fallback 缓存向量,或从 candidate._cached_vector 复用
vectors = None
if precomputed_vectors is not None:
vectors = precomputed_vectors
elif candidates and '_cached_vector' in candidates[0]:
# 从 candidate 对象上取缓存向量
vectors = [c.get('_cached_vector') for c in candidates]
if None in vectors:
vectors = None # 有 None 则退化到重新编码
if vectors is None:
try:
bge = _get_bge()
texts = [c['doc_meta']['text'] for c in candidates]
vectors = bge.encode_batch(texts, normalize_output=True)
except Exception as e:
print(f"[MMR] 向量编码失败: {e}")
return candidates[:top_k]
qv = np.array(query_vector, dtype=np.float32)
qv = qv / (np.linalg.norm(qv) + 1e-8)
@@ -453,10 +467,21 @@ def recall_memory(req: RecallRequest):
except Exception as e:
print(f"[Recall] jina rerank 失败,fallback 到 embedding 相似度排序: {e}")
try:
cand_vectors = bge.encode_batch(doc_texts, normalize_output=True)
doc_texts = [r['doc_meta'].get('text', '') for r in semantic_results]
# ══ 优化4:fallback 编码向量复用(避免 MMR 重复 encode)═
# 如果后续要走 MMR,先编码并存到 candidate 上,避免 MMR 重新编码
if req.diversity > 0 and top1_vector is not None:
fallback_vectors = bge.encode_batch(doc_texts, normalize_output=True)
for r, fv in zip(semantic_results, fallback_vectors):
r['_cached_vector'] = fv
_mmr_vectors_cache = fallback_vectors
else:
fallback_vectors = bge.encode_batch(doc_texts, normalize_output=True)
_mmr_vectors_cache = None
qv = np.array(top1_vector)
scores = []
for cv in cand_vectors:
for cv in fallback_vectors:
cv_arr = np.array(cv)
score = float(np.dot(qv, cv_arr))
scores.append(score)
+12 -4
View File
@@ -1,11 +1,15 @@
"""Jina Rerank — 模力方舟 API
API: https://ai.gitee.com/v1/rerank
模型: jina-embeddings-v4
模型: bge-reranker-v2-m3(备选:jina-reranker-v2-base-en)
用途: 搜索结果重排,提升 top-k 准确性
依赖:
pip install requests aiohttp
环境变量:
RERANK_API_KEY — 优先级最高(覆盖代码中的 key)
BGE_API_KEY — 备选(与 BGE embedder 共用同一 key)
"""
import os
import json
@@ -16,7 +20,9 @@ from pathlib import Path
# 模力方舟 API 配置
RERANK_API_URL = "https://ai.gitee.com/v1/rerank"
RERANK_API_KEY = "3TSVVXRFFECE4TISXHGE1VXDAXBIPAP6O1VPJK18"
# 优先级:环境变量 RERANK_API_KEY > BGE_API_KEY > 硬编码占位符
_RERANK_KEY_ENV = os.environ.get("RERANK_API_KEY") or os.environ.get("BGE_API_KEY") or "3TSVVXRFFECE4TISXHGE1VXDAXBIPAP6O1VPJK18"
RERANK_API_KEY = _RERANK_KEY_ENV
# 模力方舟 bge-reranker-v2-m3(不是 jina-embeddings-v4)
DEFAULT_MODEL = "bge-reranker-v2-m3"
@@ -136,12 +142,14 @@ class JinaRerankClient:
resp.raise_for_status()
result = resp.json()
# 解析结果
# 解析结果(注意:document 可能是字符串或 {text:...} 对象)
rerank_results = []
for item in result.get("results", []):
doc = item.get("document", "")
doc_text = doc.get("text", doc) if isinstance(doc, dict) else doc
r = RerankResult(
index=item["index"],
document=item["document"],
document=doc_text,
score=item["relevance_score"],
)
rerank_results.append(r)