""" 向量重排模块 实现多种重排策略以提高检索质量 支持 CrossEncoder 深度重排 """ import torch import numpy as np from typing import List, Tuple, Dict from sklearn.feature_extraction.text import TfidfVectorizer from sklearn.metrics.pairwise import cosine_similarity import re from collections import Counter import math # CrossEncoder support try: from sentence_transformers import CrossEncoder as SentenceTransformerCrossEncoder CROSSENCODER_AVAILABLE = True except ImportError: CROSSENCODER_AVAILABLE = False print("⚠️ sentence-transformers not available. CrossEncoder reranking disabled.") class DocumentReranker: """文档重排器基类""" def __init__(self): self.name = "BaseReranker" def rerank(self, query: str, documents: List[dict], top_k: int = 5) -> List[Tuple[dict, float]]: """重排文档并返回top_k结果""" raise NotImplementedError class TFIDFReranker(DocumentReranker): """基于TF-IDF的重排器""" def __init__(self): super().__init__() self.name = "TFIDFReranker" # 移除 stop_words 以支持中文，使用 char_wb 分词器 self.vectorizer = TfidfVectorizer( analyzer='char_wb', # 字符级分词，支持中文 ngram_range=(2, 4), # 2-4 字符 n-gram max_features=5000 ) def rerank(self, query: str, documents: List[dict], top_k: int = 5) -> List[Tuple[dict, float]]: """使用TF-IDF重新排序文档""" if not documents: return [] # 提取文档内容 doc_texts = [doc.page_content if hasattr(doc, 'page_content') else str(doc) for doc in documents] all_texts = [query] + doc_texts # 计算TF-IDF矩阵 tfidf_matrix = self.vectorizer.fit_transform(all_texts) query_vec = tfidf_matrix[0] doc_vecs = tfidf_matrix[1:] # 计算相似度 similarities = cosine_similarity(query_vec, doc_vecs).flatten() # 排序并返回top_k ranked_indices = np.argsort(similarities)[::-1] results = [] for i in ranked_indices[:top_k]: results.append((documents[i], float(similarities[i]))) return results class BM25Reranker(DocumentReranker): """基于BM25算法的重排器""" def __init__(self, k1: float = 1.5, b: float = 0.75): super().__init__() self.name = "BM25Reranker" self.k1 = k1 self.b = b def _tokenize(self, text: str) -> List[str]: """ 改进的分词，支持中英文 中文使用字符级分词，英文使用单词分词 """ # 检测是否包含中文 has_chinese = any('\u4e00' <= char <= '\u9fff' for char in text) if has_chinese: # 中文：使用字符级 + 2-gram chars = list(text.lower()) # 生成 unigram 和 bigram tokens = chars + [chars[i] + chars[i+1] for i in range(len(chars)-1)] return [t for t in tokens if t.strip()] # 移除空格 else: # 英文：使用单词分词 return re.findall(r'\b\w+\b', text.lower()) def _compute_idf(self, documents: List[str], query_terms: List[str]) -> Dict[str, float]: """计算IDF值""" N = len(documents) idf = {} for term in query_terms: df = sum(1 for doc in documents if term in self._tokenize(doc)) idf[term] = math.log((N - df + 0.5) / (df + 0.5)) return idf def _bm25_score(self, query_terms: List[str], document: str, avg_doc_len: float, idf: Dict[str, float]) -> float: """计算BM25分数""" doc_terms = self._tokenize(document) doc_len = len(doc_terms) term_freq = Counter(doc_terms) score = 0.0 for term in query_terms: if term in term_freq: tf = term_freq[term] score += idf.get(term, 0) * (tf * (self.k1 + 1)) / ( tf + self.k1 * (1 - self.b + self.b * doc_len / avg_doc_len) ) return score def rerank(self, query: str, documents: List[dict], top_k: int = 5) -> List[Tuple[dict, float]]: """使用BM25重新排序文档""" if not documents: return [] query_terms = self._tokenize(query) doc_texts = [doc.page_content if hasattr(doc, 'page_content') else str(doc) for doc in documents] # 计算平均文档长度 avg_doc_len = sum(len(self._tokenize(doc)) for doc in doc_texts) / len(doc_texts) # 计算IDF idf = self._compute_idf(doc_texts, query_terms) # 计算BM25分数 scores = [] for doc_text in doc_texts: score = self._bm25_score(query_terms, doc_text, avg_doc_len, idf) scores.append(score) # 排序并返回top_k ranked_indices = np.argsort(scores)[::-1] results = [] for i in ranked_indices[:top_k]: results.append((documents[i], float(scores[i]))) return results class SemanticReranker(DocumentReranker): """基于语义相似度的重排器""" def __init__(self, embeddings_model): super().__init__() self.name = "SemanticReranker" self.embeddings_model = embeddings_model def rerank(self, query: str, documents: List[dict], top_k: int = 5) -> List[Tuple[dict, float]]: """使用语义相似度重新排序文档""" if not documents: return [] # 获取查询嵌入 query_embedding = self.embeddings_model.embed_query(query) # 获取文档嵌入 doc_texts = [doc.page_content if hasattr(doc, 'page_content') else str(doc) for doc in documents] doc_embeddings = self.embeddings_model.embed_documents(doc_texts) # 计算余弦相似度 similarities = [] for doc_emb in doc_embeddings: sim = cosine_similarity([query_embedding], [doc_emb])[0][0] similarities.append(sim) # 排序并返回top_k ranked_indices = np.argsort(similarities)[::-1] results = [] for i in ranked_indices[:top_k]: results.append((documents[i], float(similarities[i]))) return results class CrossEncoderReranker(DocumentReranker): """ 基于 CrossEncoder 的重排器 使用联合编码，相比 Bi-Encoder 准确率提升 15-20% 适合精排阶段 (Top 20-100 文档) """ def __init__(self, model_name: str = "cross-encoder/ms-marco-MiniLM-L-6-v2", max_length: int = 512): """ 初始化 CrossEncoder 重排器 Args: model_name: 模型名称，默认使用轻量级模型 - "cross-encoder/ms-marco-MiniLM-L-6-v2" (轻量级，推荐) - "cross-encoder/ms-marco-MiniLM-L-12-v2" (平衡) - "BAAI/bge-reranker-base" (中文优化) - "BAAI/bge-reranker-large" (高精度) max_length: 最大输入长度 """ super().__init__() self.name = "CrossEncoderReranker" self.model_name = model_name self.max_length = max_length # 加载模型 if not CROSSENCODER_AVAILABLE: raise ImportError( "CrossEncoder requires sentence-transformers. " "Install with: pip install sentence-transformers" ) try: print(f"🔧 加载 CrossEncoder 模型: {model_name}...") self.model = SentenceTransformerCrossEncoder(model_name, max_length=max_length) print(f"✅ CrossEncoder 模型加载成功") except Exception as e: print(f"❌ CrossEncoder 模型加载失败: {e}") raise def rerank(self, query: str, documents: List[dict], top_k: int = 5) -> List[Tuple[dict, float]]: """ 使用 CrossEncoder 重新排序文档 Args: query: 查询文本 documents: 候选文档列表 top_k: 返回结果数量 Returns: 排序后的 (document, score) 元组列表 """ if not documents: return [] # 提取文档内容 doc_texts = [doc.page_content if hasattr(doc, 'page_content') else str(doc) for doc in documents] # 构造 [query, doc] 对 query_doc_pairs = [[query, doc_text] for doc_text in doc_texts] # CrossEncoder 评分 - 联合编码 try: scores = self.model.predict(query_doc_pairs) # 排序 ranked_indices = np.argsort(scores)[::-1] # 返回 top_k 结果 results = [] for i in ranked_indices[:top_k]: results.append((documents[i], float(scores[i]))) return results except Exception as e: print(f"⚠️ CrossEncoder 重排失败: {e}") # 回退到原始顺序 return [(doc, 0.0) for doc in documents[:top_k]] class HybridReranker(DocumentReranker): """混合重排器，融合多种策略""" def __init__(self, embeddings_model, weights: Dict[str, float] = None): super().__init__() self.name = "HybridReranker" # 初始化各种重排器 self.tfidf_reranker = TFIDFReranker() self.bm25_reranker = BM25Reranker() self.semantic_reranker = SemanticReranker(embeddings_model) # 设置权重 self.weights = weights or { 'tfidf': 0.3, 'bm25': 0.3, 'semantic': 0.4 } def rerank(self, query: str, documents: List[dict], top_k: int = 5) -> List[Tuple[dict, float]]: """使用混合策略重新排序文档""" if not documents: return [] # 获取各种重排结果 tfidf_results = self.tfidf_reranker.rerank(query, documents, len(documents)) bm25_results = self.bm25_reranker.rerank(query, documents, len(documents)) semantic_results = self.semantic_reranker.rerank(query, documents, len(documents)) # 创建文档到分数的映射 doc_scores = {} for doc in documents: doc_id = id(doc) doc_scores[doc_id] = {'doc': doc, 'tfidf': 0, 'bm25': 0, 'semantic': 0} # 填充各种分数 for doc, score in tfidf_results: doc_scores[id(doc)]['tfidf'] = score for doc, score in bm25_results: doc_scores[id(doc)]['bm25'] = score for doc, score in semantic_results: doc_scores[id(doc)]['semantic'] = score # 归一化分数 for score_type in ['tfidf', 'bm25', 'semantic']: scores = [info[score_type] for info in doc_scores.values()] if max(scores) > 0: max_score = max(scores) for doc_id in doc_scores: doc_scores[doc_id][score_type] /= max_score # 计算综合分数 final_scores = [] for doc_id, info in doc_scores.items(): combined_score = ( self.weights['tfidf'] * info['tfidf'] + self.weights['bm25'] * info['bm25'] + self.weights['semantic'] * info['semantic'] ) final_scores.append((info['doc'], combined_score)) # 排序并返回top_k final_scores.sort(key=lambda x: x[1], reverse=True) return final_scores[:top_k] class DiversityReranker(DocumentReranker): """多样性重排器，避免结果重复""" def __init__(self, embeddings_model, diversity_lambda: float = 0.5): super().__init__() self.name = "DiversityReranker" self.embeddings_model = embeddings_model self.diversity_lambda = diversity_lambda def _calculate_diversity_penalty(self, candidate_doc: str, selected_docs: List[str]) -> float: """计算多样性惩罚""" if not selected_docs: return 0.0 candidate_emb = self.embeddings_model.embed_documents([candidate_doc])[0] selected_embs = self.embeddings_model.embed_documents(selected_docs) max_similarity = 0.0 for selected_emb in selected_embs: sim = cosine_similarity([candidate_emb], [selected_emb])[0][0] max_similarity = max(max_similarity, sim) return max_similarity def rerank(self, query: str, documents: List[dict], top_k: int = 5) -> List[Tuple[dict, float]]: """使用多样性策略重新排序文档""" if not documents: return [] # 首先使用语义相似度获取初始排序 semantic_results = SemanticReranker(self.embeddings_model).rerank( query, documents, len(documents) ) # MMR (Maximal Marginal Relevance) 算法 selected_docs = [] selected_texts = [] remaining_docs = [doc for doc, _ in semantic_results] relevance_scores = {id(doc): score for doc, score in semantic_results} while len(selected_docs) < top_k and remaining_docs: best_score = -1 best_doc = None best_idx = -1 for i, doc in enumerate(remaining_docs): doc_text = doc.page_content if hasattr(doc, 'page_content') else str(doc) relevance = relevance_scores[id(doc)] diversity_penalty = self._calculate_diversity_penalty(doc_text, selected_texts) # MMR分数 = λ * 相关性 - (1-λ) * 多样性惩罚 mmr_score = ( self.diversity_lambda * relevance - (1 - self.diversity_lambda) * diversity_penalty ) if mmr_score > best_score: best_score = mmr_score best_doc = doc best_idx = i if best_doc is not None: selected_docs.append((best_doc, best_score)) selected_texts.append( best_doc.page_content if hasattr(best_doc, 'page_content') else str(best_doc) ) remaining_docs.pop(best_idx) return selected_docs def create_reranker(reranker_type: str, embeddings_model=None, **kwargs) -> DocumentReranker: """ 工厂函数：创建指定类型的重排器 Args: reranker_type: 重排器类型 - 'tfidf': TF-IDF 重排 - 'bm25': BM25 重排 - 'semantic': Bi-Encoder 语义重排 - 'crossencoder': CrossEncoder 重排 (推荐) ⭐ - 'hybrid': 混合重排 - 'diversity': 多样性重排 embeddings_model: 嵌入模型 (某些重排器需要) **kwargs: 其他参数 - model_name: CrossEncoder 模型名称 - max_length: CrossEncoder 最大长度 - weights: 混合重排权重 Returns: DocumentReranker: 重排器实例 """ if reranker_type.lower() == 'tfidf': return TFIDFReranker() elif reranker_type.lower() == 'bm25': return BM25Reranker(**kwargs) elif reranker_type.lower() == 'semantic': if embeddings_model is None: raise ValueError("SemanticReranker requires embeddings_model") return SemanticReranker(embeddings_model) elif reranker_type.lower() in ['crossencoder', 'cross_encoder', 'cross-encoder']: # CrossEncoder 不需要 embeddings_model，使用自己的模型 model_name = kwargs.get('model_name', 'cross-encoder/ms-marco-MiniLM-L-6-v2') max_length = kwargs.get('max_length', 512) return CrossEncoderReranker(model_name=model_name, max_length=max_length) elif reranker_type.lower() == 'hybrid': if embeddings_model is None: raise ValueError("HybridReranker requires embeddings_model") return HybridReranker(embeddings_model, **kwargs) elif reranker_type.lower() == 'diversity': if embeddings_model is None: raise ValueError("DiversityReranker requires embeddings_model") return DiversityReranker(embeddings_model, **kwargs) else: raise ValueError( f"Unknown reranker type: {reranker_type}. " f"Available types: tfidf, bm25, semantic, crossencoder, hybrid, diversity" ) # 使用示例 if __name__ == "__main__": # 模拟文档 class MockDoc: def __init__(self, content): self.page_content = content docs = [ MockDoc("人工智能是计算机科学的一个分支"), MockDoc("机器学习是人工智能的子领域"), MockDoc("深度学习使用神经网络"), MockDoc("自然语言处理处理文本数据"), MockDoc("今天天气很好") ] query = "什么是人工智能？" # 测试TF-IDF重排 tfidf_reranker = TFIDFReranker() results = tfidf_reranker.rerank(query, docs, top_k=3) print("TF-IDF重排结果:") for doc, score in results: print(f"分数: {score:.4f} - 内容: {doc.page_content}")