如何将NER模型输出的实体字典正确追加到目标字典的实体组值中？

原NER模型代码及输出tokens #

模型调用代码 #

token_classifier = pipeline('ner', model='./fine_tune_nerbert_output/', tokenizer = './fine_tune_nerbert_output/', aggregation_strategy="average")
sentence = "alisa brown i live in san diego, california and sometimes in kansas city, missouri"
tokens = token_classifier(sentence)

输出的tokens列表 #

[{'entity_group': 'LABEL_1',
  'score': 0.99938214,
  'word': 'alisa',
  'start': 0,
  'end': 5},
 {'entity_group': 'LABEL_2',
  'score': 0.9972813,
  'word': 'brown',
  'start': 6,
  'end': 11},
 {'entity_group': 'LABEL_0',
  'score': 0.99798816,
  'word': 'i live in',
  'start': 12,
  'end': 21},
 {'entity_group': 'LABEL_3',
  'score': 0.9993938,
  'word': 'san',
  'start': 22,
  'end': 25},
 {'entity_group': 'LABEL_4',
  'score': 0.9988097,
  'word': 'diego',
  'start': 26,
  'end': 31},
 {'entity_group': 'LABEL_0',
  'score': 0.9996742,
  'word': ',',
  'start': 31,
  'end': 32},
 {'entity_group': 'LABEL_3',
  'score': 0.9985813,
  'word': 'california',
  'start': 33,
  'end': 43},
 {'entity_group': 'LABEL_0',
  'score': 0.9997311,
  'word': 'and sometimes in',
  'start': 44,
  'end': 60},
 {'entity_group': 'LABEL_3',
  'score': 0.9995384,
  'word': 'kansas',
  'start': 61,
  'end': 67},
 {'entity_group': 'LABEL_4',
  'score': 0.9988242,
  'word': 'city',
  'start': 68,
  'end': 72},
 {'entity_group': 'LABEL_0',
  'score': 0.99949193,
  'word': ',',
  'start': 72,
  'end': 73},
 {'entity_group': 'LABEL_3',
  'score': 0.99960154,
  'word': 'missouri',
  'start': 74,
  'end': 82}]

当前处理代码及问题 #

现有处理循环 #

ner_dict = dict()
nested_dict = dict()
for token in tokens:
    if token['entity_group'] != 'LABEL_0':
        if token['entity_group'] in ner_dict:
            nested_dict[token['entity_group']] = {}
            nested_dict[token['entity_group']][token['word']] = token['score']
            ner_dict.update({token['entity_group']: (ner_dict[token['entity_group']], nested_dict[token['entity_group']])})
        else:
            ner_dict[token['entity_group']] = {}
            ner_dict[token['entity_group']][token['word']] = token['score']

当前错误输出（嵌套元组） #

{'LABEL_1': {'devyn': 0.9995816},
 'LABEL_2': {'donahue': 0.9996502},
 'LABEL_3': ((({'san': 0.9994766}, {'california': 0.998961}),
   {'san': 0.99925905}),
  {'california': 0.9987863}),
 'LABEL_4': ({'francisco': 0.99923646}, {'diego': 0.9992399})}

理想输出（扁平元组） #

{'LABEL_1': {'devyn': 0.9995816},
 'LABEL_2': {'donahue': 0.9996502},
 'LABEL_3': ({'san': 0.9994766}, {'california': 0.998961}, {'san': 0.99925905},
  {'california': 0.9987863}),
 'LABEL_4': ({'francisco': 0.99923646}, {'diego': 0.9992399})}

修改后的代码 #

核心思路是用列表暂存每个entity_group对应的{word:score}字典，最后统一转成元组，避免每次拼接产生嵌套：

ner_dict = dict()
for token in tokens:
    if token['entity_group'] != 'LABEL_0':
        entity = token['entity_group']
        # 生成当前token对应的字典项
        item = {token['word']: token['score']}
        if entity in ner_dict:
            # 已存在则追加到列表
            ner_dict[entity].append(item)
        else:
            # 不存在则初始化列表
            ner_dict[entity] = [item]

# 把所有列表转成元组，符合理想输出的结构
for key in ner_dict:
    ner_dict[key] = tuple(ner_dict[key])

代码说明 #

1. 初始化时，每个entity_group对应一个列表，而非单个字典，方便后续追加元素
2. 遍历token时，直接将新的{word:score}字典追加到对应列表中，不会产生嵌套
3. 最后统一把列表转为元组，匹配理想输出的不可变结构

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