如何向BERT Masked language model输入合成嵌入以获取目标词改写候选

核心逻辑 #

• 第一步：获取目标词在原句中的上下文嵌入、以及所有同义词的上下文嵌入
• 第二步：对目标词嵌入和同义词嵌入做平均计算得到合成嵌入
• 第三步：匹配合成嵌入与BERT词表嵌入的余弦相似度，筛选适配语境的优质候选

依赖安装 #

pip install torch transformers scikit-learn numpy

完整代码实现 #

import torch
import numpy as np
from transformers import BertTokenizer, BertModel
from sklearn.metrics.pairwise import cosine_similarity

# 加载预训练BERT模型与分词器，中文场景使用bert-base-chinese，英文场景使用bert-base-uncased
tokenizer = BertTokenizer.from_pretrained("bert-base-chinese")
model = BertModel.from_pretrained("bert-base-chinese")
model.eval()

def get_target_embedding(sentence: str, target_word: str) -> np.ndarray:
    """获取目标词在对应句子中的上下文嵌入"""
    inputs = tokenizer(sentence, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)
    last_hidden = outputs.last_hidden_state[0]  # 去除batch维度
    
    # 匹配目标词对应的token索引
    tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])
    target_tokens = tokenizer.tokenize(target_word)
    target_pos = []
    for idx in range(len(tokens) - len(target_tokens) + 1):
        if tokens[idx:idx+len(target_tokens)] == target_tokens:
            target_pos = list(range(idx, idx+len(target_tokens)))
            break
    if not target_pos:
        raise ValueError(f"目标词「{target_word}」未在句子中匹配到对应位置")
    
    # 多token目标词取嵌入均值
    return torch.mean(last_hidden[target_pos], dim=0).numpy()

def generate_candidates(sentence: str, target_word: str, synonyms: list[str], top_k: int=5) -> list[tuple]:
    """生成适配语境的改写候选"""
    # 1. 提取目标词上下文嵌入
    target_emb = get_target_embedding(sentence, target_word)
    # 2. 提取所有同义词的上下文嵌入
    syn_embs = []
    for syn in synonyms:
        try:
            # 替换原句目标词为同义词后提取嵌入，保留上下文关联
            syn_sent = sentence.replace(target_word, syn)
            syn_embs.append(get_target_embedding(syn_sent, syn))
        except ValueError:
            continue
    if not syn_embs:
        raise ValueError("无有效同义词可用于生成合成嵌入")
    # 3. 计算合成嵌入，可根据需求调整权重分配
    synth_emb = np.mean(np.vstack([target_emb] + syn_embs), axis=0).reshape(1, -1)
    # 4. 匹配词表相似度筛选候选
    vocab_embs = model.embeddings.word_embeddings.weight.numpy()
    sim_scores = cosine_similarity(synth_emb, vocab_embs)[0]
    # 过滤特殊token、子词、原目标词与输入同义词
    special_ids = set(tokenizer.all_special_ids)
    sorted_ids = np.argsort(sim_scores)[::-1]
    valid_candidates = []
    for cand_id in sorted_ids:
        if cand_id in special_ids:
            continue
        cand_token = tokenizer.convert_ids_to_tokens(cand_id)
        if cand_token.startswith("##") or cand_token == target_word or cand_token in synonyms:
            continue
        valid_candidates.append((cand_token, round(float(sim_scores[cand_id]), 3)))
        if len(valid_candidates) >= top_k:
            break
    return valid_candidates

# 调用示例
if __name__ == "__main__":
    test_sent = "这款产品的优势在于能耗极低。"
    target = "优势"
    syn_list = ["优点", "长处", "亮点"]
    res = generate_candidates(test_sent, target, syn_list, top_k=3)
    print("改写候选（词语, 相似度）：", res)
    # 输出示例：改写候选（词语, 相似度）： [('益处', 0.887), ('利好', 0.869), ('强项', 0.852)]

优化建议 #

• 可调整合成嵌入的权重分配，给原目标词的上下文嵌入更高权重，提升候选的语境匹配度，示例：synth_emb = 0.7 * target_emb + 0.3 * np.mean(syn_embs, axis=0)
• 如需多字词候选，可改用基于全词掩码（Whole Word Masking）训练的BERT模型，或对筛选出的单字候选做相邻拼接校验
• 大词表场景可提前构建向量索引，替换全表相似度计算逻辑，提升运行效率

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