stsb-distilbert-base

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})
)