4. 向量化嵌入模型选择

任务目标

在中文RAG场景中，不同嵌入模型对检索质量影响显著。本案例将使用同一中文问答数据集，分别通过 text-embedding-3-small、Cohere embed-multilingual-v3.0以及HuggingFace的BAAI/bge-small-zh-v1.5生成向量，基于余弦相似度进行Top-K检索，以Recall@K和命中率（Hit Rate）为指标评估三者性能差异。

环境与依赖

pip install  cohere sentence-transformers numpy pandas

需要提前设置环境变量_与COHERE_。HuggingFace模型首次运行会自动下载权重（约100MB）。

数据集准备

使用公开中文问答数据集clue/clue（取自CLUE benchmark的AFQMC子集），取前200条作为语料库，另取20条不重复的问题作为查询集。每条数据包含sentence1（查询）与sentence2（相关文档）。

完整实现代码

import os
import numpy as np
import pandas as pd
from datasets import load_dataset
from  import 
import cohere
from sentence_transformers import SentenceTransformer

# ---------- 1. 加载数据 ----------
ds = load_dataset("clue", "afqmc", split="train")
df = ds.to_pandas().sample(frac=1, random_state=42).reset_index(drop=True)
corpus = df["sentence2"].head(200).tolist()   # 语料库
queries = df["sentence1"].iloc[200:220].tolist()  # 20条查询
true_docs = df["sentence2"].iloc[200:220].tolist()  # 每条查询对应的相关文档

# ---------- 2. 定义嵌入函数 ----------
_client = (=os.environ["_"])
cohere_client = cohere.Client(os.environ["COHERE_"])
hf_model = SentenceTransformer("BAAI/bge-small-zh-v1.5")

def get_embeddings(texts, provider):
    """统一接口：输入文本列表，返回归一化向量矩阵"""
    if provider == "":
        resp = _client.embeddings.create(
            model="text-embedding-3-small", input=texts)
        vecs = [d.embedding for d in resp.data]
    elif provider == "cohere":
        resp = cohere_client.embed(texts=texts,
                                   model="embed-multilingual-v3.0",
                                   input_type="search_document")
        vecs = resp.embeddings
    elif provider == "hf":
        vecs = hf_model.encode(texts, normalize_embeddings=True)
    return np.array(vecs)

# ---------- 3. 检索评估 ----------
def evaluate(provider):
    # 生成向量（/Cohere返回后需手动归一化）
    doc_vecs = get_embeddings(corpus, provider)
    query_vecs = get_embeddings(queries, provider)
    if provider != "hf":
        doc_vecs = doc_vecs / np.linalg.norm(doc_vecs, axis=1, keepdims=True)
        query_vecs = query_vecs / np.linalg.norm(query_vecs, axis=1, keepdims=True)

    # 余弦相似度矩阵
    sim_matrix = query_vecs @ doc_vecs.T

    hit_at_1, hit_at_5, recall_at_5 = 0, 0, 0
    for i in range(len(queries)):
        top_indices = np.argsort(sim_matrix[i])[::-1][:5]
        top_docs = [corpus[idx] for idx in top_indices]
        if true_docs[i] in top_docs[:1]:
            hit_at_1 += 1
        if true_docs[i] in top_docs:
            hit_at_5 += 1
            recall_at_5 += 1  # 每个查询仅1个相关文档，Recall@5=命中率

    n = len(queries)
    return {
        "Hit@1": hit_at_1 / n,
        "Hit@5": hit_at_5 / n,
        "Recall@5": recall_at_5 / n
    }

# ---------- 4. 运行对比 ----------
if __name__ == "__main__":
    results = {}
    for provider in ["", "cohere", "hf"]:
        print(f"正在评估 {provider} ...")
        results[provider] = evaluate(provider)

    # 输出结果表
    df_result = pd.DataFrame(results).T
    df_result.columns = ["Hit@1", "Hit@5", "Recall@5"]
    print("\n===== 中文检索性能对比 =====")
    print(df_result.round(4))

运行方式

export _="sk-..."
export COHERE_="..."
python evaluate_embeddings.py

验证方法与判定标准

程序运行后会在终端打印三行结果，每行包含Hit@1、Hit@5和Recall@5三个指标。由于三个模型使用完全相同的查询与语料，可直接横向比较数值：

• 判定通过条件：三个模型均成功输出0~1之间的浮点数，且Hit@5值明显高于随机基线（随机检索期望Hit@5≈5/200=0.025），说明检索有效。
• 预期观察：通常BAAI/bge-small-zh-v1.5在中文数据上Recall@5接近或超过与Cohere（该模型针对中文优化），但具体数值取决于数据子集采样，以实际输出为准。

常见问题

• /Cohere限流：若请求过多报429错误，可在get_embeddings中增加time.sleep(0.5)或减小数据集规模。
• HuggingFace下载慢：首次运行需下载模型，可设置HF_ENDPOINT=https://hf-mirror.com加速。
• 结果差异：若某模型所有指标为0，检查是否因无效或网络问题导致嵌入返回空值。