Update app.py

app.py

@@ -2,9 +2,11 @@
 #------------------------------------------------------------------------------
 #------------------------------------------------------------------------------
 
+from aem import customroot
 import streamlit as st 
 import spacy 
 import spacy_streamlit
+from spacy import displacy
 from io import StringIO
 import pandas as pd
 
@@ -132,33 +134,64 @@ with st.spinner("Text wird eingelesen..."):
         text = st.text_area(" ", DEFAULT_TEXT, height=200)
 st.success("Text ist eingelesen!")
 
+st.markdown("---")
 #------------------------------------------------------------------------------
 #------------------------------------------------------------------------------
 
+st.markdown("### Named Entities")
 #Farben für die verschiedenen Entitäten 
 colors = {"PER": "#fdec3e", "PERSON": "#fdec3e", "LOC": "#7e56c2", "ORT": "#7e56c2", "ORG": "#209485" , "ORGANISATION": "#209485" , "MISC": "#eb4034",  "ZEIT": "#4c9c4b", "OBJEKT": "#7e56c2"}
 
 #Spacy-Streamlit NER Visualizer
-
 #NER-Prozess wird gestartet, je nach Model werden hier die entsprechenden Entitäten gewechselt.
 with st.spinner('Named Entities werden gesucht...'):
     doc = nlp(text)
     if model == "de_fnhd_nerdh":    
-        spacy_streamlit.visualize_ner(doc, labels = ["PERSON", "ORT", "ORGANISATION", "OBJEKT", "ZEIT",], show_table=False,    colors = colors)
+        entities = st.multiselect('Entitäten auswählen', ['PERSON', 'ORT', 'ORGANISATION', 'OBJEKT', 'ZEIT', 'Alle Entitäten'], default= ['Alle Entitäten'])
+        if 'Alle Entitäten' in entities:
+            entities = ['PERSON', 'ORT', 'ORGANISATION', 'OBJEKT', 'ZEIT']
+
+        options = {"ents": entities,"colors": colors}
+        ent_html = displacy.render(doc, style="ent", options=options, jupyter=False)
+        st.markdown(ent_html, unsafe_allow_html=True)
+        #spacy_streamlit.visualize_ner(doc, labels = ["PERSON", "ORT", "ORGANISATION", "OBJEKT", "ZEIT",], show_table=False,    colors = colors)
     else: 
-        spacy_streamlit.visualize_ner(doc, labels = ["PER", "LOC", "ORG", "MISC"], show_table=False,    colors = colors)
+        entities = st.multiselect('Entitäten auswählen', ["PER", "LOC", "ORG", "MISC", 'Alle Entitäten'], default= ['Alle Entitäten'])
+        if 'Alle Entitäten' in entities:
+            entities = ["PER", "LOC", "ORG", "MISC"]
+
+        options = {"ents": entities,"colors": colors}
+        ent_html = displacy.render(doc, style="ent", options=options, jupyter=False)
+        st.markdown(ent_html, unsafe_allow_html=True)
+        #spacy_streamlit.visualize_ner(doc, labels = ["PER", "LOC", "ORG", "MISC"], show_table=False,    colors = colors)
+st.markdown(' ')
 st.success('Suchprozess ist abgeschlossen!')
 
 #------------------------------------------------------------------------------
 #------------------------------------------------------------------------------
+#Download-Funktion der Entitäten
+st.sidebar.markdown('\n\n')
+st.sidebar.markdown('''
+### NER-Ergebnnisse in einer .csv-Datei downloaden.
+Die Datei enthält die ausgewählten Entitäten.
+''')
 
 #Um die NER-Ergebnisse downloaden zu können, werden die Entitäten in einer csv gespeichert
 results = []
 for ent in doc.ents:
-    results.append([ent.text,ent.label_])
+    if ent.label_ in entities:
+        results.append([ent.text,ent.label_])
 df_results = pd.DataFrame(results, columns = ['text', 'label'])
 csv = convert_df(df_results)
 
+st.sidebar.download_button( 
+    "Ergebnisse downloaden",
+   csv,
+   "ner_results_" + model + ".csv",
+   "text/csv",
+   key='download-csv'
+)
+
 
 #------------------------------------------------------------------------------
 #------------------------------------------------------------------------------
@@ -177,22 +210,6 @@ if model == "de_fnhd_nerdh":
     ```
     ''')
 
-#------------------------------------------------------------------------------
-#------------------------------------------------------------------------------
-
-#Download-Button 
-st.sidebar.markdown('\n\n')
-st.sidebar.markdown('''
-### NER-Ergebnnisse in einer .csv-Datei downloaden.
-Die Datei enthält alle Entitäts-Typen.
-''')
-st.sidebar.download_button( 
-    "Ergebnisse downloaden",
-   csv,
-   "ner_results_" + model + ".csv",
-   "text/csv",
-   key='download-csv'
-)
 
 #------------------------------------------------------------------------------
 #------------------------------------------------------------------------------