Highlighting Component Addition

app.py

@@ -49,8 +49,16 @@ def ner_tagging(text_1):
 	prediction = np.argmax(output, axis=-1)[0]
 
 	prediction = [mapping[str(i)] for i in prediction]
-
-	return prediction
+	
+	text_2 = text_1.split(" ")
+	output = []
+  	for w in range(len(text_2)):
+  		if prediction[w] != "O":
+  			output.extend([(text_2[w], prediction[w]), (" ", None)])
+  		else:
+  			output.extend([(text_2[w], None), (" ", None)])
+	
+	return output
 
 text_1 = gr.inputs.Textbox(lines=5)
 ner_tag = gr.outputs.Textbox()
@@ -61,7 +69,7 @@ with open("vocab.json",'r') as f:
 lookup_layer = keras.layers.StringLookup(vocabulary=vocab['tokens'])
 
 iface = gr.Interface(ner_tagging, 
-	inputs=text_1,outputs=ner_tag, examples=[['EU rejects German call to boycott British lamb .'], 
+	inputs=text_1,outputs=['highlight'], examples=[['EU rejects German call to boycott British lamb .'], 
 	["He said further scientific study was required and if it was found that action was needed it should be taken by the European Union ."]], title="Named Entity Recognition with Transformers",
 	description = "Named Entity Recognition with Transformers on CoNLL2003 Dataset",
         article = "Author: <a href=\"https://huggingface.co/reichenbach\">Rishav Chandra Varma</a>")