stsb-distilbert-base 项目介绍
stsb-distilbert-base 是一个由 sentence-transformers 团队开发的模型,主要用于将句子和段落映射到一个768维的密集向量空间。这项技术可以用于如聚类或语义搜索等任务。这个模型构建在 DistilBERT 的基础上,利用了其简化和高效的架构。
使用方法
要使用这个模型,用户需要安装 sentence-transformers
库。有两种主要的使用方式:一种是通过 sentence-transformers 库,另一种是通过 HuggingFace Transformers 库。
使用 sentence-transformers 库
首先,使用命令安装库:
pip install -U sentence-transformers
然后,您可以像下面这样使用该模型:
from sentence_transformers import SentenceTransformer
sentences = ["这是一个例句", "每个句子都会被转换"]
model = SentenceTransformer('sentence-transformers/stsb-distilbert-base')
embeddings = model.encode(sentences)
print(embeddings)
使用 HuggingFace Transformers 库
即使不使用 sentence-transformers 库,用户也可以通过 HuggingFace Transformers 库来使用模型。这需要对输入进行预处理,然后应用正确的池化操作以获得句子嵌入。
from transformers import AutoTokenizer, AutoModel
import torch
# 均值池化 - 结合注意力掩码以正确平均
def mean_pooling(model_output, attention_mask):
token_embeddings = model_output[0]
input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)
sentences = ['这是一个例句', '每个句子都会被转换']
# 从 HuggingFace Hub 加载模型
tokenizer = AutoTokenizer.from_pretrained('sentence-transformers/stsb-distilbert-base')
model = AutoModel.from_pretrained('sentence-transformers/stsb-distilbert-base')
# 对句子进行标记化
encoded_input = tokenizer(sentences, padding=True, truncation=True, return_tensors='pt')
# 计算标记嵌入
with torch.no_grad():
model_output = model(**encoded_input)
# 进行池化操作,这里采用均值池化
sentence_embeddings = mean_pooling(model_output, encoded_input['attention_mask'])
print("句子嵌入:")
print(sentence_embeddings)
模型评价结果
对于该模型的自动评估,用户可以参考 Sentence Embeddings Benchmark。有关具体的评估结果,可以访问:Sentence Embeddings Benchmark
完整的模型架构
该模型由以下架构组成:
SentenceTransformer(
(0): Transformer({'max_seq_length': 128, 'do_lower_case': False}) with Transformer model: DistilBertModel
(1): Pooling({'word_embedding_dimension': 768, 'pooling_mode_cls_token': False, 'pooling_mode_mean_tokens': True, 'pooling_mode_max_tokens': False, 'pooling_mode_mean_sqrt_len_tokens': False})
)
引用与作者
这个模型是由 sentence-transformers 团队训练的。如果你觉得这个模型有所帮助,可以引用他们的研究发表:Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks。