reranker.py

"""
向量重排模块
实现多种重排策略以提高检索质量
"""

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


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"
        self.vectorizer = TfidfVectorizer(stop_words='english', 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]:
        """简单分词"""
        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 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:
    """工厂函数：创建指定类型的重排器"""
    
    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() == '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}")


# 使用示例
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}")