Update app.py

app.py

@@ -100,6 +100,21 @@ with col4:
 
 runtext =st.text_area(inputNote, str(original_text2), height=300)
 
+# Extract words associated with each entity
+def genEntities(ann, entity):
+              # entity colour dict
+              #ent_col = {'DISEASE':'#B42D1B', 'CHEMICAL':'#F06292'}
+              ent_col = {'DISEASE':'pink', 'CHEMICAL':'orange'}
+              # separate into the different entities
+              entities = trans_df['Class'].unique()
+
+              if entity in entities:
+                             ent = list(trans_df[trans_df['Class']==entity]['Entity'].unique())
+                             entlist = ",".join(ent)
+                             st.markdown(f'<p style="color:{ent_col[entity]};font-size:20px;">{entlist}</p>', unsafe_allow_html=True)
+                             #for i in ent:
+                                           #st.markdown(f'<p style="color:{ent_col[entity]};font-size:20px;">{i}</p>', unsafe_allow_html=True)
+                                          
 
 def visualize (run_text,output):
     text =''
@@ -170,8 +185,12 @@ with col1:
         table["Entity"].append(key)
         table["Class"].append(ent_bc[key])
     trans_df = pd.DataFrame(table)
-    st.table(trans_df)    
-
+    st.subheader("DISEASE")
+    genEntities(trans_df, 'DISEASE')
+    st.subheader("CHEMICAL")
+    genEntities(trans_df, 'CHEMICAL')
+    #st.table(trans_df)    
+    
 with col2:
     st.button('NER')
 #    doc = nlp(str(original_text2))