【Python】sentence_transformers 库：生成句子、段落或图像的高质量嵌入（embeddings）

本文链接：https://blog.csdn.net/u013172930/article/details/147537927

sentence_transformers 是一个用于生成句子、段落或图像的高质量嵌入（embeddings）的 Python 库，基于 Hugging Face 的 transformers 库。它通过预训练的 Transformer 模型（如 BERT、RoBERTa、DistilBERT 等）生成固定长度的密集向量表示，广泛应用于自然语言处理（NLP）任务，如语义搜索、文本相似性比较、聚类、释义挖掘等。以下是对 sentence_transformers 库的详细说明。

1. sentence_transformers 库的作用

生成嵌入：将文本（句子、段落）或图像编码为固定长度的向量（嵌入），表示语义信息。
高效任务支持：支持语义文本相似性（STS）、语义搜索、聚类、问答检索、图像搜索等任务。
多语言支持：提供超过 100 种语言的预训练模型，适合多语言应用。
易于微调：允许用户基于特定任务微调模型，生成任务特定的嵌入。
跨模态支持：支持文本和图像嵌入到同一向量空间（如 CLIP 模型）。

2. 安装与环境要求

根据官方文档，推荐以下环境：

Python 版本：3.9 或更高。
依赖库：
- PyTorch 1.11.0+（支持 GPU 需安装对应 CUDA 版本）。
- transformers 4.34.0+。
安装命令：
```
pip install sentence-transformers
```
可选扩展：
- ONNX/OpenVINO：用于推理优化和量化。
- 训练支持：pip install sentence-transformers[train]
- GPU 支持：确保安装与 CUDA 版本匹配的 PyTorch。

注意：若需 GPU 加速，需根据系统 CUDA 版本安装 PyTorch，参考 PyTorch 官方指南。

3. 核心功能与用法

sentence_transformers 提供两种主要模型类型：SentenceTransformer（用于生成嵌入）和 CrossEncoder（用于重新排序或相似性评分）。

3.1 SentenceTransformer：生成嵌入

SentenceTransformer 将文本编码为固定长度的向量，适合语义搜索、相似性比较等任务。

基本用法：

from sentence_transformers import SentenceTransformer

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

# 待编码的句子
sentences = [
    "The weather is lovely today.",
    "It's so sunny outside!",
    "He drove to the stadium."
]

# 生成嵌入
embeddings = model.encode(sentences)
print(embeddings.shape)  # 输出: (3, 384)  # 3 个句子，每个嵌入 384 维

# 计算相似性
similarities = model.similarity(embeddings, embeddings)
print(similarities)
# 输出: tensor([[1.0000, 0.6660, 0.1046],
#                [0.6660, 1.0000, 0.1411],
#                [0.1046, 0.1411, 1.0000]])

说明：

模型 all-MiniLM-L6-v2 是一个轻量高效的预训练模型，生成 384 维嵌入。
model.encode() 将文本转换为嵌入（numpy 数组或 tensor）。
model.similarity() 计算嵌入之间的余弦相似度（或其他相似性度量）。

带提示（Prompt）的高级用法：
某些模型支持在推理时添加提示以优化性能，例如查询嵌入：

model = SentenceTransformer("mixedbread-ai/mxbai-embed-large-v1")
query_embedding = model.encode("What are Pandas?", prompt_name="query")
document_embeddings = model.encode([
    "Pandas is a software library written for the Python programming language for data manipulation and analysis.",
    "Pandas are a species of bear native to South Central China."
])
similarity = model.similarity(query_embedding, document_embeddings)
print(similarity)  # 输出: tensor([[0.7594, 0.7560]])

3.2 CrossEncoder：重新排序

CrossEncoder 用于对句子对进行直接评分，适合需要高精度的任务（如检索重排序）。

用法：

from sentence_transformers import CrossEncoder

# 加载预训练 CrossEncoder 模型
model = CrossEncoder("cross-encoder/ms-marco-MiniLM-L6-v2")

# 查询与文档
query = "Which planet is known as the Red Planet?"
passages = [
    "Venus is often called Earth's twin because of its similar size and proximity.",
    "Mars, known for its reddish appearance, is often referred to as the Red Planet.",
    "Jupiter, the largest planet in our solar system, has a prominent red spot."
]

# 计算评分
scores = model.predict([(query, passage) for passage in passages])
print(scores)  # 输出: [0.0123, 0.9987, 0.3456]（示例分数）

说明：

CrossEncoder 直接处理句子对，输出相似性分数（通常 0 到 1）。
适合需要高精度的场景，但计算成本高于 SentenceTransformer。

4. 预训练模型

sentence_transformers 提供超过 10,000 个预训练模型，托管在 Hugging Face Hub，覆盖以下类型：

通用嵌入模型：如 all-MiniLM-L6-v2（轻量）、all-mpnet-base-v2（高性能）。
多语言模型：如 distiluse-base-multilingual-cased-v2，支持 100+ 语言。
领域特定模型：如 multi-qa-MiniLM-L6-cos-v1（问答）、paraphrase-MiniLM-L6-v2（释义）。
跨模态模型：如 CLIP 模型（clip-ViT-B-32），支持文本和图像嵌入。

加载模型：

model = SentenceTransformer("all-MiniLM-L6-v2")  # 从 Hugging Face 下载

本地加载：
若需离线使用，可先下载模型并保存：

model = SentenceTransformer("all-MiniLM-L6-v2")
model.save("local_model_path")
# 加载本地模型
model = SentenceTransformer("local_model_path")

5. 微调与训练

sentence_transformers 支持微调模型以适应特定任务，提供多种损失函数和训练方式。

5.1 训练示例

以下是一个微调模型的示例，用于语义文本相似性任务：

from sentence_transformers import SentenceTransformer, InputExample, losses
from torch.utils.data import DataLoader

# 加载模型
model = SentenceTransformer("all-MiniLM-L6-v2")

# 准备训练数据
train_examples = [
    InputExample(texts=["The weather is nice today.", "It's pleasant outside."], label=0.9),
    InputExample(texts=["The weather is nice today.", "He went to the park."], label=0.2)
]
train_dataloader = DataLoader(train_examples, shuffle=True, batch_size=16)

# 定义损失函数
train_loss = losses.CosineSimilarityLoss(model)

# 训练模型
model.fit(
    train_objectives=[(train_dataloader, train_loss)],
    epochs=1,
    warmup_steps=100
)

# 保存模型
model.save("finetuned_model")

说明：

InputExample 定义训练样本，包含句子对和标签（如相似性分数）。
CosineSimilarityLoss 优化嵌入之间的余弦相似度。
model.fit 执行训练，支持多任务学习。

5.2 支持的损失函数

sentence_transformers.losses 提供多种损失函数，包括：

CosineSimilarityLoss：优化句子对的余弦相似度。
TripletLoss：基于三元组（锚点、正例、负例）优化。
SoftmaxLoss：用于分类任务（如 NLI）。
MultipleNegativesRankingLoss：适合无监督或弱监督训练。

5.3 训练提示

数据集：可使用 Hugging Face 数据集（如 STS-B、SNLI）或自定义数据。
提示（Prompt）：支持在训练和推理时添加提示以提升性能。
PEFT 支持：自 v3.3.0 起支持参数高效微调（Parameter-Efficient Fine-Tuning）。
NanoBEIR 评估：用于快速评估信息检索性能。

6. 性能优化

ONNX/OpenVINO：使用 ONNX 或 OpenVINO 后端进行推理优化和量化。
```
pip install sentence-transformers[onnx]
```
CPU 加速：v3.3.0 引入 OpenVINO int8 量化，CPU 推理速度提升 4.5 倍。
GPU 加速：确保 PyTorch 支持 CUDA。
批量处理：model.encode(sentences, batch_size=32) 优化内存和速度。
缓存嵌入：预计算并存储常用文本的嵌入，减少重复计算。

7. 实际应用场景

语义搜索：

from sentence_transformers import SentenceTransformer, util

model = SentenceTransformer("multi-qa-MiniLM-L6-cos-v1")
query = "How big is London?"
passages = ["London has 9,787,426 inhabitants at the 2011 census."]
query_embedding = model.encode(query)
passage_embedding = model.encode(passages)
similarity = util.dot_score(query_embedding, passage_embedding)
print(similarity)  # 输出相似度

聚类：使用嵌入进行 K-Means 或其他聚类算法。
释义挖掘：使用 util.paraphrase_mining 查找语料库中的释义。
跨模态搜索：使用 CLIP 模型实现文本-图像搜索。
问答系统：结合 CrossEncoder 进行检索和重排序。

8. 与 LangChain 集成

sentence_transformers 常与 LangChain 结合，用于构建基于语义的应用程序（如聊天机器人、文档检索）：

from sentence_transformers import SentenceTransformer
from langchain.vectorstores import FAISS
from langchain.embeddings import HuggingFaceEmbeddings

# 初始化嵌入模型
embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")

# 创建向量存储
texts = ["The weather is nice today.", "It's sunny outside."]
vector_store = FAISS.from_texts(texts, embeddings)

# 搜索
query = "How's the weather?"
results = vector_store.similarity_search(query, k=2)
print(results)

9. 注意事项

依赖体积：sentence_transformers 依赖 PyTorch 和 transformers，安装可能占用较大空间（约 5-7GB，含模型）。可使用 --no-cache-dir 减小 Docker 镜像体积。
模型下载：首次加载模型会自动从 Hugging Face 下载，需确保网络连接或预先下载模型。
离线使用：保存模型到本地路径以支持无网络环境。
版本兼容性：定期检查 sentence-transformers 的版本更新（如 v3.3.0 引入了多项优化）。
多语言模型：选择多语言模型时，需引用相关论文以确保学术规范。

10. 综合示例

以下是一个综合示例，结合嵌入生成、相似性计算和 CrossEncoder 重排序：

from sentence_transformers import SentenceTransformer, CrossEncoder, util

# 加载模型
encoder = SentenceTransformer("all-MiniLM-L6-v2")
reranker = CrossEncoder("cross-encoder/ms-marco-MiniLM-L6-v2")

# 语料库和查询
query = "What is pandas?"
corpus = [
    "Pandas is a Python library for data analysis.",
    "Pandas are animals native to China.",
    "NumPy is a Python library for numerical computing."
]

# 生成嵌入并计算相似度
query_embedding = encoder.encode(query)
corpus_embeddings = encoder.encode(corpus)
similarities = util.cos_sim(query_embedding, corpus_embeddings)

# 获取初步检索结果
top_k = 3
hits = [{"corpus_id": i, "score": similarities[0][i]} for i in range(len(corpus))]
hits = sorted(hits, key=lambda x: x["score"], reverse=True)[:top_k]

# 使用 CrossEncoder 重排序
pairs = [(query, corpus[hit["corpus_id"]]) for hit in hits]
scores = reranker.predict(pairs)

# 输出结果
for hit, score in zip(hits, scores):
    print(f"Score: {score:.4f}, Text: {corpus[hit['corpus_id']]}")

输出示例：

Score: 0.9987, Text: Pandas is a Python library for data analysis.
Score: 0.3456, Text: Pandas are animals native to China.
Score: 0.0123, Text: NumPy is a Python library for numerical computing.

11. 资源与文档

官方文档：https://www.sbert.net/
GitHub 仓库：https://github.com/UKPLab/sentence-transformers
Hugging Face 模型：https://huggingface.co/models?library=sentence-transformers
论文引用：
- SBERT: Reimers, N., & Gurevych, I. (2019). Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks.🔗
- 多语言模型: Reimers, N., & Gurevych, I. (2020). Making Monolingual Sentence Embeddings Multilingual using Knowledge Distillation.🔗