SuperAPI AI 工具集

sentence-transformers

智能体与模型

查看项目

更新于 10-15 18:37

技能介绍

用途

sentence-transformers 是一个基于 transformers 的 Python 框架,用于生成句子和文本嵌入。适用场景包括:为 RAG 提供高质量嵌入、语义相似度与搜索、文本聚类与分类、多语言嵌入(100+ 语言)、本地运行嵌入(无需 API),以及作为 OpenAI 嵌入的成本效益替代方案。

主要能力

  • 提供 5000+ 预训练模型,基于 PyTorch/Transformers。
  • 支持语义搜索、相似度计算、批量编码与模型微调。
  • 支持多语言与特定领域(法律、科学论文、代码)模型。
  • 可与 LangChain、LlamaIndex 集成。

使用方式

安装:

pip install sentence-transformers

基础用法(加载模型、生成嵌入、计算余弦相似度):

from sentence_transformers import SentenceTransformer

model = SentenceTransformer('all-MiniLM-L6-v2')

sentences = [
    "This is an example sentence",
    "Each sentence is converted to a vector"
]

embeddings = model.encode(sentences)
print(embeddings.shape)  # (2, 384)

from sentence_transformers.util import cos_sim
similarity = cos_sim(embeddings[0], embeddings[1])
print(f"Similarity: {similarity.item():.4f}")

常用模型:

  • 通用:all-MiniLM-L6-v2(384 维,快)、all-mpnet-base-v2(768 维)、all-roberta-large-v1(1024 维,较慢)。
  • 多语言:paraphrase-multilingual-MiniLM-L12-v2(50+ 语言)、paraphrase-multilingual-mpnet-base-v2(100+ 语言)。
  • 领域特定:nlpaueb/legal-bert-base-uncased(法律)、allenai/specter(科学论文)、microsoft/codebert-base(代码)。

语义搜索示例:

from sentence_transformers import SentenceTransformer, util

model = SentenceTransformer('all-MiniLM-L6-v2')

corpus = [
    "Python is a programming language",
    "Machine learning uses algorithms",
    "Neural networks are powerful"
]

corpus_embeddings = model.encode(corpus, convert_to_tensor=True)

query = "What is Python?"
query_embedding = model.encode(query, convert_to_tensor=True)

hits = util.semantic_search(query_embedding, corpus_embeddings, top_k=3)
print(hits)

相似度计算:

similarity = util.cos_sim(embedding1, embedding2)
similarity = util.dot_score(embedding1, embedding2)
similarities = util.cos_sim(embeddings, embeddings)

批量编码:

sentences = ["sentence 1", "sentence 2", ...] * 1000

embeddings = model.encode(
    sentences,
    batch_size=32,
    show_progress_bar=True,
    convert_to_tensor=False
)

微调:

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

train_examples = [
    InputExample(texts=['sentence 1', 'sentence 2'], label=0.8),
    InputExample(texts=['sentence 3', 'sentence 4'], label=0.3),
]

train_dataloader = DataLoader(train_examples, batch_size=16)
train_loss = losses.CosineSimilarityLoss(model)

model.fit(
    train_objectives=[(train_dataloader, train_loss)],
    epochs=10,
    warmup_steps=100
)

model.save('my-finetuned-model')

LangChain 集成:

from langchain_community.embeddings import HuggingFaceEmbeddings

embeddings = HuggingFaceEmbeddings(
    model_name="sentence-transformers/all-mpnet-base-v2"
)

from langchain_chroma import Chroma

vectorstore = Chroma.from_documents(
    documents=docs,
    embedding=embeddings
)

LlamaIndex 集成:

from llama_index.embeddings.huggingface import HuggingFaceEmbedding

embed_model = HuggingFaceEmbedding(
    model_name="sentence-transformers/all-mpnet-base-v2"
)

from llama_index.core import Settings
Settings.embed_model = embed_model

index = VectorStoreIndex.from_documents(documents)

适用场景

  • RAG 嵌入生成、语义相似度与搜索、文本聚类与分类、多语言任务、本地嵌入计算。
  • 可替代方案:OpenAI Embeddings(API 方式,质量最高)、Instructor(任务特定指令)、Cohere Embed(托管服务)。

模型选择指南

模型 维度 速度 质量 用途
all-MiniLM-L6-v2 384 快 良好 通用、原型
all-mpnet-base-v2 768 中 较好 生产 RAG
all-roberta-large-v1 1024 慢 最佳 高精度需求
paraphrase-multilingual 768 中 良好 多语言

限制与最佳实践

最佳实践:从 all-MiniLM-L6-v2 开始作为基线;归一化嵌入以改善余弦相似度;如有 GPU 可提速约 10 倍;使用批量编码;缓存嵌入以避免重复计算;针对领域微调以提升质量;测试不同模型(质量因任务而异);监控内存(大模型需更多 RAM)。

性能参考:MiniLM 约 2000 句/秒、120MB、384 维;MPNet 约 600 句/秒、420MB、768 维;RoBERTa 约 300 句/秒、1.3GB、1024 维。

资源:GitHub https://github.com/UKPLab/sentence-transformers(15,700+ 星标);模型 https://huggingface.co/sentence-transformers;文档 https://www.sbert.net;许可 Apache 2.0。