Update app.py

app.py

@@ -96,8 +96,16 @@ if Run_Button == True:
         else:
             spacy_display["ents"].append({"start": entity["start"], "end": entity["end"], "label": entity["entity"]})
     
-    entity_list = ["PER", "LOC", "ORG", "MISC", "person", "location", "geopolitical area", "organization", "event", "group", "date", "ordinal number", "cardinal number", "product", "quantity", "time", "money", "percent", "language", "product"]
+    tner_entity_list = ["person", "group", "facility", "organization", "geopolitical area", "location", "product", "event", "work of art", "law", "language", "date", "time", "percent", "money", "quantity", "ordinal number", "cardinal number"]
+    spacy_entity_list = ["PERSON", "NORP", "FAC", "ORG", "GPE", "LOC", "PRODUCT", "EVENT", "WORK_OF_ART", "LAW", "LANGUAGE", "DATE", "TIME", "PERCENT", "MONEY", "QUANTITY", "ORDINAL", "CARDINAL", "MISC"]
+    
+    for ent in spacy_display["ents"]:
+        if model_checkpoint == "tner/tner-xlm-roberta-base-ontonotes5":
+            ent["label"] = spacy_entity_list[tner_entity_list.index(ent["label"])]
+        else:
+            if ent["label"] == "PER": ent["label"] = "PERSON"
+    
     # colors = {'PER': '#85DCDF', 'LOC': '#DF85DC', 'ORG': '#DCDF85', 'MISC': '#85ABDF',}
-    html = spacy.displacy.render(spacy_display, style="ent", minify=True, manual=True, options={"ents": entity_list}) # , "colors": colors})
+    html = spacy.displacy.render(spacy_display, style="ent", minify=True, manual=True, options={"ents": spacy_entity_list}) # , "colors": colors})
     style = "<style>mark.entity { display: inline-block }</style>"
     st.write(f"{style}{get_html(html)}", unsafe_allow_html=True)