Update app.py

app.py

@@ -40,6 +40,39 @@ def init_ner_pipeline():
     pipe = pipeline("ner", model=model, tokenizer=tokenizer, aggregation_strategy="simple") # pass device=0 if using gpu
     return pipe
 
+
+def get_formatted_text_for_annotation(output):
+    colour_map = {'Sex': '#5DD75D',
+     'Duration': '#D92E45',
+     'Sign_symptom': '#793F41',
+     'Frequency': '#232AE7',
+     'Detailed_description': '#E1D8D1',
+     'History': '#296FB8',
+     'Clinical_event': '#E840A7',
+     'Lab_value': '#FE90C3',
+     'Age': '#31404C',
+     'Biological_structure': '#1A4B5B',
+     'Diagnostic_procedure': '#804E7A'}
+    
+    annotated_texts = []
+    next_index = 0
+    for entity in output:
+        if entity['start'] == next_index:
+    #         print("found entity")
+            extracted_text = text[entity['start']:entity['end']]
+    #         print("annotated",annotated_text)
+            annotated_texts.append((extracted_text ,entity['entity_group'],colour_map[entity['entity_group']]))
+        else:
+            unannotated_text = text[next_index:entity['start']-1]
+            annotated_texts.append(unannotated_text)
+            extracted_text = text[entity['start']:entity['end']]
+            annotated_texts.append((extracted_text ,entity['entity_group'],colour_map[entity['entity_group']]))
+            next_index =entity['end'] +1
+    
+    if next_index < len(text):
+        annotated_texts.append(text[next_index-1:len(text)-1])
+        
+    return tuple(annotated_texts)
     
 # Model initialization    
 pipeline_summarization = init_text_summarization_model()
@@ -69,6 +102,8 @@ elif selected_menu == "Summarize Document":
 elif selected_menu == "Extract Entities":
     text = get_text_from_ocr_engine()
     output = pipeline_ner (text)
+    entities_text =get_formatted_text_for_annotation(output)
+    annotated_text(entities_text)
 
 
 elif selected_menu == "Get Answers":