Update app.py

app.py

@@ -9,6 +9,10 @@ login(token = api_key)
 
 
 get_completion = pipeline("ner", model="elnasharomar2/PUNCERT_single_stages_50_epochs")
+label_names = ["O","QE","EX","QM","DOT","COM","SEMICOL","COL"]
+label_symbol = ["O",'؟!','!','؟','.','،','؛',':']
+
+id2label = {i: label for i, label in zip(label_names,label_symbol)}
 
 def merge_tokens(tokens):
     merged_tokens = []
@@ -28,7 +32,17 @@ def merge_tokens(tokens):
 def ner(input):
     output = get_completion(input)
     merged_tokens = merge_tokens(output)
-    return {"text": input, "entities": merged_tokens}
+    ## modification
+    result =""
+    idx = 0
+    for i in output:
+        result += text[idx:i["start"]]
+        result += text[i["start"]:i["end"]] + f"{id2label[i["entity_group"]]}"
+        idx = i["end"]
+    result+=text[idx:]
+    print(result)
+    return {"text": result}
+    # return {"text": input, "entities": merged_tokens}
 
 gr.close_all()
 demo = gr.Interface(fn=ner,