混合搜索架构设计：侦探的两种查案绝技

import mathimport refrom collections import Counter class SparseBM25:    """稀疏检索：BM25 算法 + 倒排索引"""     def __init__(self, k1=1.5, b=0.75):        self.k1 = k1        self.b = b        self.doc_count = 0        self.avgdl = 0.0        self.docs = []        self.doc_lens = []        self.df = Counter()        self.inverted_index = {}     def train(self, docs):        self.docs = docs        self.doc_count = len(docs)        total_len = 0        for doc in docs:            terms = self._tokenize(doc)            self.doc_lens.append(len(terms))            total_len += len(terms)            unique_terms = set(terms)            for term in unique_terms:                self.df[term] += 1                if term not in self.inverted_index:                    self.inverted_index[term] = []                self.inverted_index[term].append(self.doc_count - len(self.docs) + len(self.docs) - len(self.inverted_index[term]) - 1)        # 修正倒排索引构建，保证 ID 正确        self.inverted_index = {}        for idx, doc in enumerate(docs):            terms = set(self._tokenize(doc))            for t in terms:                self.inverted_index.setdefault(t, []).append(idx)        self.avgdl = total_len / self.doc_count     def _tokenize(self, text):        return re.findall(r"[\w]+", text.lower())     def score(self, query):        scores = [0.0] * self.doc_count        query_terms = self._tokenize(query)        for idx, doc in enumerate(self.docs):            term_counts = Counter(self._tokenize(doc))            total = 0.0            for term in query_terms:                if term in term_counts:                    df_t = self.df.get(term, 0)                    idf = math.log((self.doc_count - df_t + 0.5) / (df_t + 0.5) + 1)                    tf = term_counts[term]                    dl = self.doc_lens[idx]                    denom = tf + self.k1 * (1 - self.b + self.b * dl / self.avgdl)                    total += idf * (tf * (self.k1 + 1)) / denom            scores[idx] = total        return scores  class DenseRetriever:    """稠密检索：用随机向量模拟语义向量（实际项目用深度学习模型生成）"""     def __init__(self, dim=64):        self.dim = dim        self.doc_vectors = []        self.docs = []     def train(self, docs):        self.docs = docs        for doc in docs:            vec = self._embed(doc)            self.doc_vectors.append(vec)     def _embed(self, text):        # 模拟：用词频向量代替真实的语义向量，仅用于演示        terms = re.findall(r"[\w]+", text.lower())        vec = [0.0] * self.dim        for i, t in enumerate(terms[:self.dim]):            vec[i] = hash(t) % 100 / 100.0        return vec     def _cosine(self, v1, v2):        dot = sum(a * b for a, b in zip(v1, v2))        n1 = math.sqrt(sum(a * a for a in v1))        n2 = math.sqrt(sum(b * b for b in v2))        if n1 == 0 or n2 == 0:            return 0.0        return dot / (n1 * n2)     def score(self, query):        q_vec = self._embed(query)        return [self._cosine(q_vec, dv) for dv in self.doc_vectors]  class HybridSearch:    """混合搜索：融合稀疏与稠密两路召回结果"""     def __init__(self, alpha=0.7, k=60):        self.alpha = alpha        self.k = k        self.sparse = SparseBM25()        self.dense = DenseRetriever()     def train(self, docs):        self.sparse.train(docs)        self.dense.train(docs)     def fuse_rrf(self, ranking_list, k=60):        """RRF 融合：按排名倒数累加得分"""        score_map = {}        for ranking in ranking_list:            for rank, doc_id in enumerate(ranking, start=1):                score_map[doc_id] = score_map.get(doc_id, 0) + 1.0 / (k + rank)        return score_map     def fuse_weighted(self, sparse_scores, dense_scores):        """加权融合：归一化后按 alpha 加权相加"""        s_max = max(sparse_scores)        d_max = max(dense_scores)        fused = []        for s, d in zip(sparse_scores, dense_scores):            s_norm = s / s_max if s_max > 0 else 0            d_norm = d / d_max if d_max > 0 else 0            fused.append(self.alpha * d_norm + (1 - self.alpha) * s_norm)        return fused     def search(self, query, strategy="rrf"):        sparse_scores = self.sparse.score(query)        dense_scores = self.dense.score(query)        if strategy == "rrf":            s_ranking = [i for i, v in sorted(enumerate(sparse_scores), key=lambda x: x[1], reverse=True)]            d_ranking = [i for i, v in sorted(enumerate(dense_scores), key=lambda x: x[1], reverse=True)]            fused = self.fuse_rrf([s_ranking, d_ranking])        elif strategy == "weighted":            fused = self.fuse_weighted(sparse_scores, dense_scores)        else:            raise ValueError("Unknown strategy")        results = sorted(fused.items(), key=lambda x: x[1], reverse=True)        return [(idx, self.dense.docs[idx], score) for idx, score in results]  if __name__ == "__main__":    docs = [        "苹果发布新款手机 iPhone 15",        "macbook pro 笔记本今日开售",        "iPhone 15 价格与配置解析",        "苹果公司发布第四季度财报",        "红色连衣裙女夏款 2025 新款",        "小黑裙女装显瘦气质连衣裙",    ]    hs = HybridSearch(alpha=0.7)    hs.train(docs)    print("查询: iPhone 15 价格")    for idx, doc, score in hs.search("iPhone 15 价格", strategy="rrf"):        print(f"  rank={idx} score={score:.4f} doc={doc}")    print("查询: 黑色裙子 女")    for idx, doc, score in hs.search("黑色裙子 女", strategy="weighted"):        print(f"  rank={idx} score={score:.4f} doc={doc}")