Update app.py

app.py

@@ -1,21 +1,23 @@
-import gradio as gr
-from transformers import pipeline
-import os
-from dotenv import load_dotenv
-load_dotenv()
-
-share = os.getenv("SHARE", False)
-
-pipe = pipeline("token-classification", model="bnsapa/cybersecurity-ner")
-def tag(input):
-    k = pipe(input)
-    return f"{k}"
-
-iface = gr.Interface(fn=tag, inputs="text", outputs="text", title="Cybersecurity NER", description="Named Entity Recognition for Cybersecurity")
-
-if __name__ == "__main__":
-    if share:
-        server = "0.0.0.0"
-    else:
-        server = "127.0.0.1"
+import gradio as gr
+from transformers import pipeline
+import os
+from dotenv import load_dotenv
+load_dotenv()
+
+share = os.getenv("SHARE", False)
+
+pipe = pipeline("token-classification", model="bnsapa/cybersecurity-ner")
+def tag(input):
+    k = pipe(input)
+    for i in k:
+        input = input.replace(i["word"], i["word"] + "[" + i["entity"] + " Score: " + str(i["score"]) +"]")
+    return input
+
+iface = gr.Interface(fn=tag, inputs="text", outputs="text", title="Cybersecurity NER", description="Named Entity Recognition for Cybersecurity")
+
+if __name__ == "__main__":
+    if share:
+        server = "0.0.0.0"
+    else:
+        server = "127.0.0.1"
     iface.launch(server_name = server)