Update app.py

app.py

@@ -28,8 +28,11 @@ def predict(username, pwd, label, payload_text = None):
     # 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, prediction[0]['label'], prediction[0]['score'] * 100
+    scores = {
+        prediction[0]['label']: prediction[0]['score'] * 100
+
+    }
+    return payload_text, scores
 
 input_elements = [gr.Textbox(label="Enter Username"), gr.Textbox(label="Enter Password"), gr.Dropdown(["Malicious", "Benign"], label="Expected", info="Enter expected value"),
             gr.Textbox(label="Enter Payload", info="Optional if username and password entered already")]
@@ -39,7 +42,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"), gr.Slider(0, 100, label="Score")]
+    outputs=[gr.Textbox(label="Generated Payload"), gr.Label(label="Scores")]
 )
 demo.launch(debug=True)
 # gr.Interface.load("models/ankush-003/nosqli_identifier").launch()