Embedding模型选型实战:BGE vs OpenAI vs m3e的真实对比

2023-05-10 | Embedding 评测 RAG

为公司的知识库检索系统选型Embedding模型,我在真实业务数据上测试了8个模型。公开benchmark得分高的模型,在我们的场景不一定好。这篇文章记录完整的评测方法和结论。

为什么不能只看MTEB榜单

MTEB是目前最权威的Embedding评测基准,但它有几个问题:

数据分布不同:MTEB是通用数据集,你的业务数据分布可能完全不同
任务类型不同:MTEB包含分类、聚类等任务,但RAG主要需要检索能力
中文覆盖有限:中文评测集相对较少

评测方法

构建测试集

import json
from typing import List, Dict

def build_eval_dataset(
    docs: List[str],
    queries: List[str],
    relevance: Dict[str, List[str]]  # query_id -> [relevant_doc_ids]
) -> Dict:
    """
    构建检索评测数据集
    
    我们的数据来源:
    1. 从搜索日志抽取500个真实query
    2. 人工标注每个query的相关文档(1-5篇)
    """
    return {
        "corpus": {str(i): doc for i, doc in enumerate(docs)},
        "queries": {str(i): q for i, q in enumerate(queries)},
        "relevance": relevance
    }

# 加载业务数据
eval_data = build_eval_dataset(
    docs=load_knowledge_base(),      # 5000篇技术文档
    queries=load_search_logs(500),   # 500个真实query
    relevance=load_annotations()     # 人工标注
)

评测指标

import numpy as np
from typing import List

def recall_at_k(retrieved: List[str], relevant: List[str], k: int) -> float:
    """Recall@K: 前K个结果中包含了多少相关文档"""
    retrieved_k = set(retrieved[:k])
    relevant_set = set(relevant)
    return len(retrieved_k & relevant_set) / len(relevant_set) if relevant_set else 0

def mrr(retrieved: List[str], relevant: List[str]) -> float:
    """MRR: 第一个相关文档的排名倒数"""
    relevant_set = set(relevant)
    for i, doc_id in enumerate(retrieved):
        if doc_id in relevant_set:
            return 1.0 / (i + 1)
    return 0.0

def ndcg_at_k(retrieved: List[str], relevant: List[str], k: int) -> float:
    """NDCG@K: 考虑位置的相关性评分"""
    dcg = 0.0
    relevant_set = set(relevant)
    for i, doc_id in enumerate(retrieved[:k]):
        if doc_id in relevant_set:
            dcg += 1.0 / np.log2(i + 2)
    
    # Ideal DCG
    idcg = sum(1.0 / np.log2(i + 2) for i in range(min(len(relevant), k)))
    return dcg / idcg if idcg > 0 else 0

候选模型

模型	维度	参数量	中文支持	开源
text-embedding-ada-002	1536	?	一般	❌
text-embedding-3-small	1536	?	好	❌
BGE-large-zh-v1.5	1024	326M	优秀	✅
BGE-m3	1024	568M	优秀	✅
m3e-base	768	102M	好	✅
text2vec-large-chinese	1024	326M	好	✅

评测结果

检索效果

模型	Recall@5	Recall@10	MRR	NDCG@10
BGE-large-zh-v1.5	78.2%	85.6%	0.721	0.752
BGE-m3	76.8%	84.3%	0.705	0.738
text-embedding-3-small	72.1%	81.5%	0.682	0.715
text-embedding-ada-002	68.5%	78.2%	0.645	0.683
m3e-base	65.3%	75.8%	0.612	0.658
text2vec-large-chinese	62.1%	73.4%	0.589	0.635

⚠️ 意外发现: OpenAI的ada-002在我们的中文技术文档上表现不如BGE。这和MTEB榜单的排名不同,说明一定要在自己的数据上测试。

性能与成本

模型	推理速度(doc/s)	成本(100万doc)	部署方式
BGE-large-zh-v1.5	~150	电费	自部署
text-embedding-ada-002	~500	$100	API
m3e-base	~300	电费	自部署

BGE使用指南

from FlagEmbedding import FlagModel

# 加载模型
model = FlagModel(
    'BAAI/bge-large-zh-v1.5',
    query_instruction_for_retrieval="为这个句子生成表示以用于检索相关文章:",
    use_fp16=True
)

# Query和Document的编码方式不同!
def encode_query(query: str):
    # Query要加instruction
    return model.encode_queries([query])[0]

def encode_docs(docs: list):
    # Document不加instruction
    return model.encode(docs)

最终选择:BGE-large-zh-v1.5

理由:

在我们的数据上效果最好
可以私有化部署,数据安全
免费开源,无API成本

更新记录:
2023-05-10: 初版发布
2023-08-15: 补充BGE-m3测试
2024-01-20: 更新text-embedding-3-small数据