Update app.py

app.py

@@ -1,17 +1,21 @@
 import gradio as gr
-def predict(payload, payload_label):
-    inputs = tokenizer(test_payload, return_tensors="tf")
+from transformers import AutoTokenizer
+
+tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")
+
+def predict(payload, malitious):
+    inputs = tokenizer(payload, return_tensors="tf")
     model = TFAutoModelForSequenceClassification.from_pretrained("models/ankush-003/nosqli_identifier")
     logits = model(**inputs).logits
     predicted_class_id = int(tf.math.argmax(logits, axis=-1)[0])
     # print(model.config.id2label[predicted_class_id])
-    return model.config.id2label[predicted_class_id]
-    # return "Fuck You"
+    expected = "Malitious" if malitious else "Benign"
+    return model.config.id2label[predicted_class_id], expected 
 
 demo = gr.Interface(
     fn=predict,
     inputs=["text","checkbox"],
-    outputs="text",
+    outputs=["text","text"]
 )
 demo.launch(debug=True)
 # gr.Interface.load("models/ankush-003/nosqli_identifier").launch()