使用Sentence Transformers计算文本嵌入向量的完整指南

使用Sentence Transformers计算文本嵌入向量的完整指南

什么是文本嵌入向量

文本嵌入向量(Embeddings)是将文本转换为固定长度的向量表示的技术。这些向量能够捕捉文本的语义信息，相似的文本在向量空间中距离更近。Sentence Transformers项目提供了高效计算文本嵌入向量的工具。

安装与基础使用

首先需要安装Sentence Transformers库。安装完成后，计算文本嵌入向量只需要简单的几个步骤：

from sentence_transformers import SentenceTransformer

# 1. 加载预训练模型
model = SentenceTransformer("all-MiniLM-L6-v2")

# 准备要编码的句子
sentences = [
    "今天天气真好",
    "外面阳光明媚！",
    "他开车去了体育场"
]

# 2. 计算嵌入向量
embeddings = model.encode(sentences)
print(embeddings.shape)  # 输出: (3, 384)

# 3. 计算相似度矩阵
similarities = model.similarity(embeddings, embeddings)
print(similarities)

输出结果会显示一个3x3的相似度矩阵，对角线为1(自我相似度)，其他值表示不同句子间的语义相似度。

模型初始化详解

Sentence Transformers支持多种预训练模型，也可以加载本地训练好的模型：

# 加载预训练模型
model = SentenceTransformer("all-mpnet-base-v2")

# 或者加载本地模型
model = SentenceTransformer("output/models/mpnet-base-finetuned-all-nli")

模型会自动使用可用的最佳计算设备(如CUDA或MPS)。也可以手动指定设备：

model = SentenceTransformer("all-mpnet-base-v2", device="cuda")

提示模板(Prompt Templates)的使用

某些模型需要特定的提示模板才能达到最佳性能。例如：

• multilingual-e5-large模型：查询前应添加"query: "，段落前添加"passage: "
• bge-large-en-v1.5模型：检索时应在输入文本前添加"Represent this sentence for searching relevant passages: "

可以在初始化模型时指定提示模板：

model = SentenceTransformer(
    "intfloat/multilingual-e5-large",
    prompts={
        "classification": "对以下文本进行分类: ",
        "retrieval": "检索语义相似的文本: ",
        "clustering": "根据文本识别主题: ",
    },
    default_prompt_name="retrieval"
)

使用时有三种方式应用提示模板：

1. 在encode方法中直接指定prompt参数
2. 使用预定义的prompt_name参数
3. 使用default_prompt_name指定的默认提示

处理长文本输入

Transformer模型(如BERT)的计算复杂度和内存需求与输入长度呈平方关系增长。每个模型都有最大序列长度限制：

print("最大序列长度:", model.max_seq_length)  # 通常为256或512

# 可以调整(但不能超过模型支持的最大值)
model.max_seq_length = 200

超过最大长度的文本会被自动截断。需要注意的是，如果模型是在短文本上训练的，对长文本的表示可能不够理想。

多进程/多GPU加速

对于大规模文本处理，可以使用多进程或多GPU加速计算。核心方法是start_multi_process_pool，它会启动多个进程用于并行编码。

实际应用建议

1. 模型选择：根据任务需求选择合适的预训练模型。小型模型速度快但精度略低，大型模型精度高但资源消耗大。

2. 批量处理：对于大量文本，建议使用批量处理而非单条处理，可以提高效率。

3. 相似度计算：得到的嵌入向量可用于计算余弦相似度、欧氏距离等，用于语义搜索、聚类等任务。

4. 领域适配：对于特定领域任务，建议在领域数据上微调模型以获得更好的表示。

通过掌握这些技巧，你可以高效地利用Sentence Transformers为各种NLP任务计算高质量的文本嵌入向量。