Update app.py

app.py

@@ -5,11 +5,15 @@ import tensorflow as tf
 # from transformers import TFAutoModelForSequenceClassification
 
 # Load model directly
-from transformers import AutoTokenizer, TFAutoModelForSequenceClassification
+# from transformers import AutoTokenizer, TFAutoModelForSequenceClassification
 
-# tokenizer = AutoTokenizer.from_pretrained("ankush-003/nosqli_identifier")
-tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")
-model = TFAutoModelForSequenceClassification.from_pretrained("ankush-003/nosqli_identifier")
+# # tokenizer = AutoTokenizer.from_pretrained("ankush-003/nosqli_identifier")
+# tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")
+# model = TFAutoModelForSequenceClassification.from_pretrained("ankush-003/nosqli_identifier")
+from transformers import pipeline
+
+classifier = pipeline("sentiment-analysis", model="ankush-003/nosqli_identifier")
+# classifier(payload)
 
 def predict(username, pwd, label, payload_text = None):
     if(payload_text is None or payload_text is ""):
@@ -18,13 +22,14 @@ def predict(username, pwd, label, payload_text = None):
             "password": pwd
         }
         payload_text = json.dumps(payload)
-    inputs = tokenizer(payload_text, return_tensors="tf")
+    # 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])
+    # logits = model(**inputs).logits
+    # predicted_class_id = int(tf.math.argmax(logits, axis=-1)[0])
     # print(model.config.id2label[predicted_class_id])
+    prediction = classifier(payload_text)    
     
-    return payload_text, model.config.id2label[predicted_class_id]
+    return payload_text, prediction['label'], prediction['score'] * 100
 
 input_elements = [gr.Textbox(label="Enter Username"), gr.Textbox(label="Enter Password"), gr.Dropdown(["Malitious", "Benign"], label="Expected", info="Enter expected value"),
             gr.Textbox(label="Enter Payload", info="Optional if username and password entered already")]
@@ -34,7 +39,7 @@ demo = gr.Interface(
     description="DistilBERT-based NoSQL Injection Payload Detection Model ",
     fn=predict,
     inputs=input_elements,
-    outputs=[gr.Textbox(label="Generated Payload"), gr.Textbox(label="Model Prediction")]
+    outputs=[gr.Textbox(label="Generated Payload"), gr.Textbox(label="Model Prediction"), gr.Slider(0, 100, label="Score")]
 )
 demo.launch(debug=True)
 # gr.Interface.load("models/ankush-003/nosqli_identifier").launch()