Python sentence-transformers库示例详解

pip install sentence-transformers

pip install sentence-transformers[torch]

from sentence_transformers import SentenceTransformer
# 加载预训练模型（首次运行会自动下载）
model = SentenceTransformer('all-MiniLM-L6-v2')
# 待编码的句子
sentences = [
    "这是一个示例句子。",
    "另一个句子用来测试。",
    "这两句话看起来不太一样。"
]
# 编码成向量（每个句子变成一个固定维度的向量，如 384 维）
embeddings = model.encode(sentences)
print(embeddings.shape)  # (3, 384)

from sklearn.metrics.pairwise import cosine_similarity
# 计算相似度矩阵
similarity_matrix = cosine_similarity(embeddings)
print(similarity_matrix)

import torch
device = 'cuda' if torch.cuda.is_available() else 'cpu'
model.to(device)
embeddings = model.encode(sentences, device=device)

from sentence_transformers import SentenceTransformer, InputExample, losses
from torch.utils.data import DataLoader
# 构建训练样本
train_examples = [
    InputExample(texts=["春天来了", "天气变暖了"], label=0.8),
    InputExample(texts=["我爱中国", "我是中国人"], label=0.9),
    ...
]
# 创建数据加载器
train_dataloader = DataLoader(train_examples, shuffle=True, batch_size=16)
# 加载模型
model = SentenceTransformer('bert-base-chinese')
# 使用余弦相似度损失函数
train_loss = losses.CosineSimilarityLoss(model)
# 开始训练
model.fit(
    train_objectives=[(train_dataloader, train_loss)],
    epochs=3,
    output_path='./my_model/'
)