STS-B 数据集详解：衡量句子之间的语义相似度

一、什么是 STS-B？

STS-B 全称是 Semantic Textual Similarity Benchmark，是由 SemEval 比赛发展而来的标准句子对语义相似度数据集，用于评估模型理解两个句子语义接近程度的能力。

属于 GLUE 基准任务之一。

二、任务目标：句子相似度回归

与分类任务不同，STS-B 是一个 回归任务，模型需输出两个句子之间的相似度得分（0~5 的浮点数）：

• 5：完全等价

• 0：完全不相关

三、数据结构与样例

每条样本包含两句话与一个相似度分数：

四、数据集规模（GLUE版本）

虽然规模不大，但挑战性高，因为它要求模型具有细腻的语义理解能力。

五、评价指标

由于是回归任务，主要评价指标是：

• Pearson correlation：线性相关度

• Spearman correlation：等级相关性

这两者衡量模型输出与人类评分的相似性。

六、与 QQP/MRPC 的区别

七、HuggingFace 示例代码

1. 加载数据

from datasets import load_dataset

dataset = load_dataset("glue", "stsb")
print(dataset["train"][0])

输出：

{
  "sentence1": "A plane is taking off.",
  "sentence2": "An aircraft is taking flight.",
  "label": 4.8
}

2. Tokenizer & 模型加载

from transformers import BertTokenizer, BertForSequenceClassification

tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
model = BertForSequenceClassification.from_pretrained(
    "bert-base-uncased",
    num_labels=1  # 回归任务
)

3. 数据预处理

def tokenize(example):
    return tokenizer(example["sentence1"], example["sentence2"], truncation=True, padding="max_length")

encoded_dataset = dataset.map(tokenize, batched=True)

4. 模型训练与评估

from transformers import Trainer, TrainingArguments

args = TrainingArguments(
    output_dir="./stsb_out",
    evaluation_strategy="epoch",
    per_device_train_batch_size=16,
    num_train_epochs=3,
)

trainer = Trainer(
    model=model,
    args=args,
    train_dataset=encoded_dataset["train"],
    eval_dataset=encoded_dataset["validation"]
)

trainer.train()

5. 预测相似度

import torch

inputs = tokenizer("A man is singing.", "A person is performing vocally.", return_tensors="pt")
outputs = model(**inputs)
similarity_score = outputs.logits.item()
print(f"相似度评分：{similarity_score:.2f}")

八、实际应用场景

九、模型推荐

十、 总结