Update app.py

app.py

@@ -31,17 +31,20 @@ elif context == "New Text":
 def setModel(model_checkpoint, aggregation):
     model = AutoModelForTokenClassification.from_pretrained(model_checkpoint)
     tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)
-    return pipeline('token-classification', model=model, tokenizer=tokenizer, aggregation_strategy=aggregation)
+    return pipeline('token-classification', model=model, tokenizer=tokenizer, aggregation_strategy=aggregation), model.config.id2label
         
 Run_Button = st.button("Run", key=None)
 if Run_Button == True:
     
-    ner_pipeline = setModel(model_checkpoint, aggregation)
+    ner_pipeline, id2label = setModel(model_checkpoint, aggregation)
     output = ner_pipeline(input_text)
 
-    print(output)
+    output_new = []
+    for o in output:
+        o["entity_group"] = id2label[o["entity_group"].split("_")[-1]]
+        output_new.append(o)
 
-    df = pd.DataFrame.from_dict(output)
+    df = pd.DataFrame.from_dict(output_new)
 
     if aggregation != "none":
         df.rename(index=str,columns={'entity_group':'POS Tag'},inplace=True)