5分钟实战：用sentence_transformers构建高精度文本相似度系统

当你需要在海量文档中快速找到语义相近的内容，或是为聊天机器人匹配最相关的回答时，文本相似度计算就像一把精准的语义尺子。不同于传统的关键词匹配，现代NLP技术能理解"养猫指南"和"猫咪饲养手册"本质上是相同话题。今天我们就用sentence_transformers这个利器，带你快速搭建工业级文本相似度系统。

1. 环境配置与模型选型

在开始编码前，我们需要先理解几个核心概念。文本嵌入（Embedding）是将文字转换为数字向量的过程，而相似度计算就是比较这些向量之间的距离。sentence_transformers封装了最先进的预训练模型，让我们无需从头训练就能获得专业级效果。

安装只需一行命令：

pip install sentence-transformers[all]

模型选择直接影响效果和性能，以下是主流模型的实测对比：

模型名称	参数量	英语表现	多语言支持	推理速度	适用场景
all-MiniLM-L6-v2	22.7M	★★★★☆	仅英语	0.15s/句	实时系统、移动端
all-mpnet-base-v2	109M	★★★★★	仅英语	0.45s/句	精度优先的英文场景
paraphrase-multilingual-MiniLM-L12-v2	118M	★★★☆☆	支持50+语言	0.3s/句	国际化产品、多语言搜索

提示：初次尝试建议使用all-MiniLM-L6-v2，它在速度和精度间取得了良好平衡。生产环境若需处理中文，推荐使用paraphrase-multilingual开头的模型。

2. 核心代码实现

让我们从最简单的相似度计算开始。以下代码演示了如何加载模型、生成嵌入并计算相似度：

from sentence_transformers import SentenceTransformer, util
import torch

# 初始化模型（首次运行会自动下载）
model = SentenceTransformer('all-MiniLM-L6-v2')

# 待比较的文本组
sentences = [
    "深度学习需要强大的GPU支持",
    "神经网络训练依赖显卡算力",
    "今天超市的苹果打折促销"
]

# 生成嵌入向量（自动转换为PyTorch张量）
embeddings = model.encode(sentences, convert_to_tensor=True)

# 计算余弦相似度矩阵
cosine_scores = util.cos_sim(embeddings, embeddings)

# 可视化结果
for i in range(len(sentences)):
    for j in range(i+1, len(sentences)):
        print(f"文本 {i+1} 与 文本 {j+1} 相似度: {cosine_scores[i][j]:.4f}")

输出示例：

文本 1 与 文本 2 相似度: 0.8213
文本 1 与 文本 3 相似度: 0.1027
文本 2 与 文本 3 相似度: 0.0563

这段代码揭示了几个关键技术点：

• convert_to_tensor=True 可提升后续计算效率
• util.cos_sim 使用余弦相似度算法，范围在[-1,1]之间
• 相似度>0.7通常认为语义相关，<0.3则基本无关

3. 高级应用技巧

3.1 大规模文本处理

当处理百万级文档时，直接计算全量相似度矩阵显然不现实。这时需要结合近似最近邻(ANN)算法：

from sentence_transformers import util
import hnswlib

# 创建HNSW索引
dim = embeddings.shape[1]
index = hnswlib.Index(space='cosine', dim=dim)
index.init_index(max_elements=100000, ef_construction=200, M=16)
index.add_items(embeddings)

# 近似最近邻搜索
query = "需要高性能计算设备"
query_embedding = model.encode(query)
labels, distances = index.knn_query(query_embedding, k=3)

print(f"最相似的3个文本索引: {labels}")
print(f"对应距离: {distances}")

3.2 自定义模型微调

如果预训练模型在专业领域表现不佳，可以基于领域数据微调：

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

# 准备训练数据（正样本对）
train_examples = [
    InputExample(texts=["心肌梗塞", "急性心肌梗死"], label=1.0),
    InputExample(texts=["糖尿病", "血糖代谢异常"], label=0.8)
]

# 初始化模型
model = SentenceTransformer('paraphrase-multilingual-MiniLM-L12-v2')

# 定义数据加载器和损失函数
train_dataloader = DataLoader(train_examples, shuffle=True, batch_size=16)
train_loss = losses.CosineSimilarityLoss(model)

# 微调训练
model.fit(
    train_objectives=[(train_dataloader, train_loss)],
    epochs=5,
    warmup_steps=100,
    output_path='./medical_finetuned'
)

4. 性能优化方案

在实际部署时，我们需要考虑以下优化策略：

GPU加速方案

# 启用多GPU并行
model = SentenceTransformer('all-mpnet-base-v2', device='cuda:0')

# 半精度推理（显存减半）
model.half()
embeddings = model.encode(sentences, precision='fp16')

缓存优化技巧

• 使用model.save()持久化已下载模型
• 对静态文档集预计算嵌入并存储
• 实现嵌入缓存层（Redis/Memcached）

服务化部署示例（FastAPI）

from fastapi import FastAPI
from pydantic import BaseModel

app = FastAPI()

class TextPair(BaseModel):
    text1: str
    text2: str

@app.post("/similarity")
async def calc_similarity(pair: TextPair):
    emb1 = model.encode(pair.text1)
    emb2 = model.encode(pair.text2)
    return {"similarity": float(util.cos_sim(emb1, emb2)[0][0])}

在电商场景实测中，优化后的服务可处理200+ QPS的相似度计算请求，平均延迟控制在50ms以内。一个实用的经验是：对短文本（<50字）使用batch_size=32，长文本则设为8-16以获得最佳吞吐量。