transformers pipeline aggregation_strategy处理非BIO输出格式的模型

from transformers import AutoModelForTokenClassification, BertTokenizerFast
from transformers import pipeline

tokenizer = BertTokenizerFast.from_pretrained('models/bert-base-chinese')
model = AutoModelForTokenClassification.from_pretrained('models/albert-base-chinese-ner')
nlp = pipeline("ner", model=model, tokenizer=tokenizer,)
example = "我叫王大，喜欢去旺角餐厅吃牛角包, 今年买了阿里巴巴的股票，亏得舅老爷的裤衩都没了，我的手机号是13587677888"
nlp(example)

输出结果

[{'entity': 'B-PERSON',
  'score': 0.9999951,
  'index': 3,
  'word': '王',
  'start': 2,
  'end': 3},
 {'entity': 'E-PERSON',
  'score': 0.9999939,
  'index': 4,
  'word': '大',
  'start': 3,
  'end': 4},
 {'entity': 'B-FAC',
  'score': 0.40783358,
  'index': 9,
  'word': '旺',
  'start': 8,
  'end': 9},
 {'entity': 'I-FAC',
  'score': 0.8651608,
  'index': 10,
  'word': '角',
  'start': 9,
  'end': 10},
 {'entity': 'I-FAC',
  'score': 0.579798,
  'index': 11,
  'word': '餐',
  'start': 10,
  'end': 11},
 {'entity': 'E-FAC',
  'score': 0.86096364,
  'index': 12,
  'word': '厅',
  'start': 11,
  'end': 12},
 {'entity': 'B-DATE',
  'score': 0.9999201,
  'index': 18,
  'word': '今',
  'start': 18,
  'end': 19},
 {'entity': 'E-DATE',
  'score': 0.99997747,
  'index': 19,
  'word': '年',
  'start': 19,
  'end': 20},
 {'entity': 'B-ORG',
  'score': 0.9996025,
  'index': 22,
  'word': '阿',
  'start': 22,
  'end': 23},
 {'entity': 'I-ORG',
  'score': 0.99899405,
  'index': 23,
  'word': '里',
  'start': 23,
  'end': 24},
 {'entity': 'I-ORG',
  'score': 0.99911004,
  'index': 24,
  'word': '巴',
  'start': 24,
  'end': 25},
 {'entity': 'E-ORG',
  'score': 0.9999074,
  'index': 25,
  'word': '巴',
  'start': 25,
  'end': 26}]

但是如果想使用simple聚合策略，会发现聚合失败

nlp = pipeline("ner", model=model, tokenizer=tokenizer, aggregation_strategy='simple')
example = "我叫王大，喜欢去旺角餐厅吃牛角包, 今年买了阿里巴巴的股票，亏得舅老爷的裤衩都没了，我的手机号是13587677n88"
nlp(example)

输出结果：

[{'entity_group': 'PERSON',
  'score': 0.9999951,
  'word': '王',
  'start': 2,
  'end': 3},
 {'entity_group': 'PERSON',
  'score': 0.9999939,
  'word': '大',
  'start': 3,
  'end': 4},
 {'entity_group': 'FAC',
  'score': 0.61759746,
  'word': '旺 角 餐',
  'start': 8,
  'end': 11},
 {'entity_group': 'FAC',
  'score': 0.86096364,
  'word': '厅',
  'start': 11,
  'end': 12},
 {'entity_group': 'DATE',
  'score': 0.9999201,
  'word': '今',
  'start': 18,
  'end': 19},
 {'entity_group': 'DATE',
  'score': 0.99997747,
  'word': '年',
  'start': 19,
  'end': 20},
 {'entity_group': 'ORG',
  'score': 0.9992356,
  'word': '阿 里 巴',
  'start': 22,
  'end': 25},
 {'entity_group': 'ORG',
  'score': 0.9999074,
  'word': '巴',
  'start': 25,
  'end': 26}]

原因就是simple聚合策略只支持BI聚合，不支持BIE聚合

修改方法

把模型输出转换成BIO输出，只需要把标签E转化成I即可

model.config.id2label = {k: v.replace('E-', 'I-') for k, v in model.config.id2label.items()}
nlp = pipeline("ner", model=model, tokenizer=tokenizer, aggregation_strategy='simple')
example = "我叫王大，喜欢去旺角餐厅吃牛角包, 今年买了阿里巴巴的股票，亏得舅老爷的裤衩都没了，我的手机号是13587677n88"
nlp(example)

这次聚合就正常了

[{'entity_group': 'PERSON',
  'score': 0.99999213,
  'word': '王 大',
  'start': 2,
  'end': 4},
 {'entity_group': 'FAC',
  'score': 0.6240196,
  'word': '旺 角 餐 厅',
  'start': 8,
  'end': 12},
 {'entity_group': 'DATE',
  'score': 0.9999214,
  'word': '今 年',
  'start': 18,
  'end': 20},
 {'entity_group': 'ORG',
  'score': 0.9991275,
  'word': '阿 里 巴 巴',
  'start': 22,
  'end': 26},
 {'entity_group': 'PERSON',
  'score': 0.6155525,
  'word': '##87677n',
  'start': 51,
  'end': 57}]