从零构建cherry搜索结果压缩与重排模型：新手避坑指南

1次阅读

共计 2587 个字符，预计需要花费 7 分钟才能阅读完成。

在搜索系统中，我们经常会遇到两个主要问题：冗余结果太多，以及排序不够精准。比如，当用户搜索 ”Python 机器学习教程 ” 时，传统搜索引擎可能会返回大量内容重复的博客，或者把过时的教程排在最前面。这不仅浪费服务器资源，还严重影响用户体验。今天我们就来聊聊如何用 cherry 模型解决这些问题。

在进入具体实现之前，我们先比较下几种常见的排序方法：

TF-IDF：简单快速，但无法理解语义相似度。比如会把 ”Python” 和 ” 蟒蛇 ” 视为完全不同的词
BM25：考虑了词频和文档长度，对短文本效果不错，但同样存在语义鸿沟问题
神经排序模型：能够理解语义关系，但计算成本较高

对于大多数现代搜索场景，我们推荐使用基于 Sentence-BERT 的方案，它在精度和性能之间取得了很好的平衡。

from sentence_transformers import SentenceTransformer
import numpy as np

# 加载预训练模型
model = SentenceTransformer('paraphrase-MiniLM-L6-v2')

def extract_features(documents):
    """
    将文本列表转换为嵌入向量
    :param documents: 文本列表
    :return: numpy 数组，每行代表一个文档的向量
    """
    return model.encode(documents, convert_to_numpy=True)

from sklearn.metrics.pairwise import cosine_similarity

def remove_duplicates(documents, vectors, threshold=0.9):
    """
    基于余弦相似度去除重复文档
    :param documents: 原始文档列表
    :param vectors: 对应的向量
    :param threshold: 相似度阈值，高于此值视为重复
    :return: 去重后的文档列表
    """
    sim_matrix = cosine_similarity(vectors)
    unique_indices = []

    for i in range(len(documents)):
        # 检查是否与已选文档高度相似
        is_unique = True
        for j in unique_indices:
            if sim_matrix[i][j] > threshold:
                is_unique = False
                break
        if is_unique:
            unique_indices.append(i)

    return [documents[i] for i in unique_indices]

def rerank_results(query, documents, vectors, top_k=10):
    """
    根据查询与文档的相关性重新排序
    :param query: 搜索查询文本
    :param documents: 待排序文档列表
    :param vectors: 文档向量
    :param top_k: 返回的结果数量
    :return: 排序后的文档列表
    """
    # 获取查询向量
    query_vec = model.encode([query])[0]

    # 计算与每个文档的相似度
    similarities = cosine_similarity([query_vec], vectors)[0]

    # 获取相似度最高的 top_k 个文档索引
    top_indices = np.argsort(similarities)[-top_k:][::-1]

    return [documents[i] for i in top_indices]

使用批处理而不是单条处理：model.encode支持传入文档列表
启用 GPU 加速：确保安装了 cudatoolkit 并设置device='cuda'

from concurrent.futures import ThreadPoolExecutor

def parallel_encode(documents, batch_size=32, workers=4):
    """
    并行编码大量文档
    :param documents: 文档列表
    :param batch_size: 每个线程处理的批量大小
    :param workers: 线程数
    :return: 向量矩阵
    """
    batches = [documents[i:i + batch_size] 
               for i in range(0, len(documents), batch_size)]

    with ThreadPoolExecutor(max_workers=workers) as executor:
        results = list(executor.map(model.encode, batches))

    return np.vstack(results)

import psutil
import os

def check_memory_usage():
    """
    检查当前进程的内存使用情况
    :return: 内存使用量(MB)
    """
    process = psutil.Process(os.getpid())
    return process.memory_info().rss / 1024 / 1024  # 转换为 MB

使用模型仓库 (如 HuggingFace Hub) 管理不同版本
为每个部署创建独立的 Docker 镜像

try:
    results = rerank_results(query, docs, vectors)
except Exception as e:
    # 记录详细错误日志
    logger.error(f"Reranking failed: {str(e)}")
    # 回退到基本排序
    results = fallback_ranking(query, docs)