NER模型推理后如何将分词token还原为原始词汇并匹配对应标签？

实现代码 #

你可以通过自定义子词合并函数实现标签与文本的映射：

def merge_subwords(tokens: list, tags: list) -> list[tuple]:
    merged = []
    current_word = ""
    current_tag = ""
    # 特殊token过滤列表，可根据实际情况扩展
    special_tokens = {"[CLS]", "[SEP]", "[PAD]", "[UNK]"}
    
    for token, tag in zip(tokens, tags):
        # 跳过特殊token
        if token in special_tokens:
            continue
        # 处理##开头的子词
        if token.startswith("##"):
            current_word += token[2:]
        else:
            # 遇到新词，先保存上一组合并结果
            if current_word:
                merged.append((current_word, current_tag))
            current_word = token
            # 子词对应的X标签直接继承首词的有效标签
            current_tag = tag if tag != "X" else current_tag
    # 补充最后一组结果
    if current_word:
        merged.append((current_word, current_tag))
    return merged

调用方式 #

在你获取到tokens和tags的位置直接调用该函数即可：

for i in range(batch['input_ids'].shape[0]):
    tags = [self.unique_labels[label_id] for label_id in output[i]]
    tokens = [t for t in self.bert_tokenizer.convert_ids_to_tokens(batch['input_ids'][i]) if t != '[PAD]']
    # 调用合并函数
    merged_result = merge_subwords(tokens, tags)
    # 输出示例：[('Amazon', 'B-COMPANY'), ('and', 'O'), ('Tesla', 'B-COMPANY'), ('are', 'O'), ...]

逻辑说明 #

• 子词对应的X标签无实际意义，统一使用实体第一个子词的标签即可
• 如果你的标签体系包含I-类的实体内标签，只需调整current_tag的赋值逻辑，适配连续实体的合并规则即可

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