Archived
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:
+36
-11
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user