def test_get_entities(self):
seq = ['B-PERSON', 'I-PERSON', 'O', 'B-LOC', 'I-LOC']
chunks = get_entities(seq)
expected_chunks = [('PERSON', 0, 2), ('LOC', 3, 5)]
self.assertEqual(chunks, expected_chunks)
seq = ['B-PERSON', 'I-PERSON', 'O', 'B-LOC', 'O']
chunks = get_entities(seq)
expected_chunks = [('PERSON', 0, 2), ('LOC', 3, 4)]
self.assertEqual(chunks, expected_chunks)
seq = ['B-PERSON', 'I-PERSON', 'O', 'I-LOC', 'O']
chunks = get_entities(seq)
expected_chunks = [('PERSON', 0, 2)]
self.assertEqual(chunks, expected_chunks)
seq = ['B-PERSON', 'I-PERSON', 'O', 'O', 'B-LOC']
chunks = get_entities(seq)
expected_chunks = [('PERSON', 0, 2), ('LOC', 4, 5)]
self.assertEqual(chunks, expected_chunks)
seq = ['O', 'B-PERSON', 'O', 'O', 'B-LOC']
chunks = get_entities(seq)
expected_chunks = [('PERSON', 1, 2), ('LOC', 4, 5)]
self.assertEqual(chunks, expected_chunks)
def _build_response(self, sent, tags, prob):
words = self.tokenizer(sent)
res = {'text': sent, 'words': words, 'entities': []}
chunks = get_entities(tags)
for chunk_type, chunk_start, chunk_end in chunks:
entity = {
'text': ' '.join(words[chunk_start:chunk_end]),
'type': chunk_type,
'score': float(np.average(prob[chunk_start:chunk_end])),
'beginOffset': chunk_start,
'endOffset': chunk_end
}
res['entities'].append(entity)
return res
def _get_chunks(self, words, tags):
"""
Args:
words: sequence of word
tags: sequence of labels
Returns:
dict of entities for a sequence
Example:
words = ['President', 'Obama', 'is', 'speaking', 'at', 'the', 'White', 'House', '.']
tags = ['O', 'B-Person', 'O', 'O', 'O', 'O', 'B-Location', 'I-Location', 'O']
result = {'Person': ['Obama'], 'LOCATION': ['White House']}
"""
chunks = get_entities(tags)
res = defaultdict(list)
for chunk_type, chunk_start, chunk_end in chunks:
res[chunk_type].append(' '.join(
words[chunk_start:chunk_end])) # todo delimiter changeable
return res