Update app.py

app.py

@@ -1,4 +1,5 @@
 import gradio as gr
+import json
 import tensorflow as tf
 # from transformers import AutoTokenizer
 # from transformers import TFAutoModelForSequenceClassification
@@ -10,8 +11,12 @@ from transformers import AutoTokenizer, TFAutoModelForSequenceClassification
 tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")
 model = TFAutoModelForSequenceClassification.from_pretrained("ankush-003/nosqli_identifier")
 
-def predict(payload, label):
-    inputs = tokenizer(payload, return_tensors="tf")
+def predict(username, pwd, payload_text = None, label):
+    if(payload_text is None):
+        payload["username"] = username
+        payload["password"] = pwd
+        payload_text = json.dumps(payload)
+    inputs = tokenizer(payload_text, return_tensors="tf")
     # model = TFAutoModelForSequenceClassification.from_pretrained("ankush-003/nosqli_identifier")
     logits = model(**inputs).logits
     predicted_class_id = int(tf.math.argmax(logits, axis=-1)[0])
@@ -21,7 +26,8 @@ def predict(payload, label):
 
 demo = gr.Interface(
     fn=predict,
-    inputs=["text",
+    inputs=[gr.Textbox(label="Enter Username"),gr.Textbox(label="Enter Password"),
+            gr.Textbox(label="Enter Payload[optional]"),
            gr.Dropdown(
             ["Malitious", "Benign"], label="Expected", info="Enter expected value"
         )],