Update app.py

app.py

@@ -1,7 +1,8 @@
 import streamlit as st
-from transformers import pipeline, AutoModelForSequenceClassification, AutoTokenizer, AutoModelForTokenClassification, AutoModelWithLMHead
+from transformers import pipeline, AutoModelForTokenClassification, AutoTokenizer
 import pandas as pd
 import spacy
+from spacy import displacy
 
 st.set_page_config(layout="wide")
 
@@ -106,17 +107,12 @@ if Run_Button and input_text != "":
     spacy_display["title"] = None
 
     for entity in output_comb:
-        spacy_display["ents"].append({"start": entity["start"], "end": entity["end"], "label": entity["entity_group"]})
-        
-    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 == "asahi417/tner-xlm-roberta-base-ontonotes5":
-            ent["label"] = spacy_entity_list[tner_entity_list.index(ent["label"])]
-        else:
-            if ent["label"] == "PER": ent["label"] = "PERSON"
-    
-    html = spacy.displacy.render(spacy_display, style="ent", minify=True, manual=True, options={"ents": spacy_entity_list})
+        label = entity["entity_group"]
+        # Make sure the label is not empty
+        if label not in ["PERSON", "NORP", "FAC", "ORG", "GPE", "LOC", "PRODUCT", "EVENT", "WORK_OF_ART", "LAW", "LANGUAGE", "DATE", "TIME", "PERCENT", "MONEY", "QUANTITY", "ORDINAL", "CARDINAL", "MISC"]:
+            label = "UNKNOWN"
+        spacy_display["ents"].append({"start": entity["start"], "end": entity["end"], "label": label})
+
+    html = displacy.render(spacy_display, style="ent", minify=True, manual=True, options={"ents": ["PERSON", "NORP", "FAC", "ORG", "GPE", "LOC", "PRODUCT", "EVENT", "WORK_OF_ART", "LAW", "LANGUAGE", "DATE", "TIME", "PERCENT", "MONEY", "QUANTITY", "ORDINAL", "CARDINAL", "MISC"]})
     style = "<style>mark.entity { display: inline-block }</style>"
     st.write(f"{style}{get_html(html)}", unsafe_allow_html=True)