from __future__ import annotations
from typing import List
from ehr import HealthRecord
from annotations import Entity
from collections import defaultdict
import re
class DictNER:
'''
A dictionary based NER model.
'''
def __init__(self):
self.ner_re: dict = {}
def _get_clean_re(self, entity_list: List[str]) -> str:
'''
Generates a regular expression from a list of entities
Parameters
----------
entity_list : List[str]
List of entity text.
Returns
-------
entity_re : str
Regular expression.
'''
regex_chars = ['(', ')', '[', ']', '{', '}', '+', '*', '?', '$', '^', '&']
for i in range(len(entity_list)):
# We need to add a \ so it does not take entity text as regex
# character
for char in regex_chars:
entity_list[i] = entity_list[i].replace(char,
'\\' + char)
# A space/new line/tab before and after the text to indicate
# a seperate word
entity_re = '[\n| |\t]|[\n| |\t]'.join(entity_list)
entity_re = '[\n| |\t]' + entity_re + '[\n| |\t]'
return entity_re
def fit(self, train_data: List[HealthRecord]) -> DictNER:
'''
Generates a dictionary for the model
Parameters
----------
train_data : List[HealthRecord]
Records to generate the dictionary from.
Returns
-------
DictNER
Self object.
'''
ner_dict = defaultdict(list)
for data in train_data:
for ent in data.entities.values():
# We have a specific RE for Dosage
if ent.name != 'Strength':
# Ignore text with length 1
if ent.ann_text.lower() not in ner_dict[ent.name]\
and len(ent.ann_text) > 1:
ner_dict[ent.name].append(ent.ann_text.lower())
for name, entity_list in ner_dict.items():
ner_dict[name] = self._get_clean_re(entity_list)
# Dosage is just a number followed by mg or mcg
ner_dict['Strength'] = '\d+[ ]*(?:mg|mcg)'
self.ner_re = dict(ner_dict)
return self
def predict(self, test_data: List[HealthRecord])\
-> List[List[Entity]]:
'''
Returns character ranges for all predicted entities
Parameters
----------
test_data : List[HealthRecord]
Text to predict the entities.
Returns
-------
List[List[Entity]]
Predictions for each example. Each prediction list
contains several Entity objects.
'''
predictions = []
for data in test_data:
entities = []
j = 1
for ent_name, ent_re in self.ner_re.items():
# Get the start and end character ranges of entities
# Remove the extra space at the start and end of entity
ranges = [(m.start(0) + 1, m.end(0) - 1, ent_name) \
for m in re.finditer(ent_re, data.text, re.IGNORECASE)]
# Convert to Entity Objects
for r in ranges:
ent = Entity(entity_id = "T" + str(j))
ent.set_range([r[0], r[1]])
ent.set_entity_type(r[2])
entities.append(ent)
j += 1
predictions.append(entities)
return predictions
