normalized diseases

utils.py

@@ -4,6 +4,7 @@ import pandas as pd
 import ssl
 import torch
 import re
+import difflib
 from pprint import pprint
 from captum.attr import visualization
 
@@ -21,6 +22,39 @@ class PyTMinMaxScalerVectorized(object):
         tensor.mul_(scale).sub_(tensor.min(dim=0, keepdim=True)[0])
         return tensor
     
+def _normalized_diseases(text_list, disease):
+    candidates = difflib.get_close_matches(disease, text_list)
+    if len(candidates) > 0:
+        return candidates[0]
+    return ''
+    
+def clean_disease_string(disease):
+    disease = disease.strip().lower()
+    disease = re.sub(r'[^\w\s]','',disease)
+    return disease
+
+def normalized_diseases(text, disease_list):
+    disease_list = list(set(disease_list))
+    text_split = text.split()
+    normalized = []
+    for disease in disease_list:
+        # case when the disease is one word
+        if ' ' not in disease:
+            candidate = _normalized_diseases(disease=disease, text_list=text_split)
+            if len(candidate) > 0:
+                candidate = clean_disease_string(candidate)
+                normalized.append(candidate)
+        else:
+            concept = ''
+            for disease_word in disease.split():
+                candidate = _normalized_diseases(text_list=text_split, disease=disease_word)
+                if len(candidate) > 0:
+                    concept += (candidate + ' ')
+            if len(concept.split()) == len(disease.split()):
+                concept = clean_disease_string(concept)
+                normalized.append(concept)
+    return list(set(normalized))
+    
 def get_diseases(text, pipe):
     results = pipe(text)
     diseases = []
@@ -44,7 +78,8 @@ def get_diseases(text, pipe):
     if len(disease_span) > 1:
         disease = text[disease_span[0]: disease_span[1]]
         diseases.append(disease)
-    return diseases    
+    normalized = normalized_diseases(text, diseases)
+    return normalized    
 
 def find_end(text):
     """Find the end of the report."""