如何用HuggingFace Transformers实现基于BERT的连续句子分类并提取[SEP]嵌入？

核心流程步骤 #

1. 构造多句子输入
   使用Transformers的tokenizer直接处理句子列表，它会自动在句子间插入[SEP]令牌，并在开头添加[CLS]，同时生成对应的input_ids、attention_mask等张量。

2. 模型前向传播获取隐藏状态
   将输入传入BERT模型，获取最后一层（或指定层）的所有令牌隐藏状态。

3. 定位并提取[SEP]嵌入
   通过input_ids找到所有[SEP]令牌的位置索引，再从隐藏状态中提取对应位置的嵌入向量。

完整代码示例 #

from transformers import AutoTokenizer, AutoModel
import torch

# 加载预训练模型和分词器
model_name = "bert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModel.from_pretrained(model_name)

# 示例多句子输入（支持批量处理）
sentence_groups = [
    ["This is sentence 1.", "Another sentence here."],
    ["First sentence of group 2.", "Second one.", "Third sentence."]
]

# 处理输入：自动插入[CLS]和[SEP]，生成张量
inputs = tokenizer(sentence_groups, return_tensors="pt", padding=True, truncation=True)

# 前向传播获取隐藏状态（推理阶段关闭梯度计算）
with torch.no_grad():
    outputs = model(**inputs)
# last_hidden_state形状：(batch_size, sequence_length, hidden_size)
last_hidden_states = outputs.last_hidden_state

# 获取[SEP]令牌的ID
sep_token_id = tokenizer.sep_token_id

# 定位所有[SEP]的位置：返回(batch_indices, seq_indices)
batch_indices, seq_indices = (inputs["input_ids"] == sep_token_id).nonzero(as_tuple=True)

# 提取所有[SEP]的嵌入
all_sep_embeddings = last_hidden_states[batch_indices, seq_indices, :]

# 可选：按样本分组提取每个样本的[SEP]嵌入
sep_embeddings_per_sample = []
batch_size = inputs["input_ids"].shape[0]
for idx in range(batch_size):
    # 找到当前样本对应的所有[SEP]位置
    sample_seq_indices = seq_indices[batch_indices == idx]
    # 提取嵌入
    sample_sep_embeds = last_hidden_states[idx, sample_seq_indices, :]
    sep_embeddings_per_sample.append(sample_sep_embeds)

# 输出示例：第一个样本的[SEP]嵌入形状
print(f"第一个样本的[SEP]嵌入形状: {sep_embeddings_per_sample[0].shape}")

关键细节说明 #

• 批量处理兼容：代码支持同时处理多个多句子样本，padding和truncation会自动处理序列长度不一致的情况。
• 多层隐藏状态可选：如果需要提取中间层的[SEP]嵌入，可使用outputs.hidden_states（需在加载模型时设置output_hidden_states=True），例如取第12层隐藏状态：hidden_states = outputs.hidden_states[12]。
• 微调场景适配：如果是在微调模型中提取嵌入，无需设置model.eval()，但要注意梯度计算的开销，可根据需求选择是否使用torch.no_grad()。

内容的提问来源于stack exchange，提问作者Tom Leung