Update app.py

app.py

@@ -42,7 +42,7 @@ model_checkpoint = st.sidebar.radio("", model_list)
 #     aggregation = st.sidebar.radio("", ('first', 'simple', 'average', 'max', 'none'))
     
 # st.sidebar.write("Please refer 'https://huggingface.co/transformers/_modules/transformers/pipelines/token_classification.html' for entity grouping with aggregation_strategy parameter.")
-aggregation = 'first'
+aggregation = 'none'
 st.subheader("Select Text Input Method")
 input_method = st.radio("", ('Select from Examples', 'Write or Paste New Text'))
 if input_method == 'Select from Examples':
@@ -89,12 +89,24 @@ if Run_Button == True:
 
     for entity in output:
         if aggregation != "none":
-            spacy_display["ents"].append({"start": entity["start"], "end": entity["end"], "label": entity["entity_group"]})
+            spacy_display["ents"].append({"start": entity["start"], "end": entity["end"], "label": entity_map[entity["entity_group"]]})
         else:
-            spacy_display["ents"].append({"start": entity["start"], "end": entity["end"], "label": entity["entity"]})
-    
-    entity_list = ["PER", "LOC", "ORG", "MISC"]
-    colors = {'PER': '#85DCDF', 'LOC': '#DF85DC', 'ORG': '#DCDF85', 'MISC': '#85ABDF',}
+            spacy_display["ents"].append({"start": entity["start"], "end": entity["end"], "label": entity_map[entity["entity"]]})
+
+    entity_map = {'LABEL_0': 'O',
+                    'LABEL_1': 'B-Intervention',
+                    'LABEL_2': 'I-Intervention',
+                    'LABEL_3': 'B-Outcome',
+                    'LABEL_4': 'I-Outcome',
+                    'LABEL_5': 'B-Values',
+                    'LABEL_6': 'I-Value'}
+
+    colors = {'B-Intervention': '#85DCDF', 
+                'I-Intervention': '#85DCDF',
+                'B-Outcome': '#DF85DC', 
+                'I-Outcome': '#DF85DC',
+                'B-Value': '#DCDF85',
+                'I-Value': '#DCDF85'}
     html = spacy.displacy.render(spacy_display, style="ent", minify=True, manual=True, options={"ents": entity_list, "colors": colors})
     style = "<style>mark.entity { display: inline-block }</style>"
     st.write(f"{style}{get_html(html)}", unsafe_allow_html=True)