INNER CODE UNIT · Python
get_ner_tags
mlwithme/BertWithPretrained · Tasks/TaskForChineseNER.py:188
def get_ner_tags(logits, token_ids, entities, SEP_IDX=102):
"""
:param logits: [src_len,batch_size,num_samples]
:param token_ids: # [src_len,batch_size]
:return:
e.g.
logits = torch.tensor([[[0.4, 0.7, 0.2],[0.5, 0.4, 0.1],[0.1, 0.2, 0.3],[0.5, 0.7, 0.2],[0.1, 0.2, 0.5]],
[[0.3, 0.2, 0.5],[0.7, 0.8, 0.4],[0.1, 0.1, 0.3],[0.9, 0.2, 0.1],[0.1, 0.5,0.2]]])
logits = logits.transpose(0, 1) # [src_len,batch_size,num_samples]
token_ids = torch.tensor([[101, 2769, 511, 102, 0],
[101, 56, 33, 22, 102]]).transpose(0, 1) # [src_len,batch_size]
labels, probs = get_ner_tags(logits, token_ids, entities)
[['O', 'B-LOC'], ['B-ORG', 'B-LOC', 'O']]
[[0.5, 0.30000001192092896], [0.800000011920929, 0.30000001192092896, 0.8999999761581421]]
"""
# entities = {'O': 0, 'B-ORG': 1, 'B-LOC': 2, 'B-PER': 3, 'I-ORG': 4, 'I-LOC': 5, 'I-PER': 6}
label_list = list(entities.keys())
logits = logits[1:].transpose(0, 1) # [batch_size,src_len-1,num_samples]