Adding Probability

app.py

@@ -45,7 +45,9 @@ def predict(newdf, loaded_model):
     
     prediction = loaded_model.predict(newdf)
 
-    return prediction 
+    proba = loaded_model.predict_proba(newdf)
+
+    return prediction, proba[0] 
     
 def get_label(newpred):
     if newpred == 0:
@@ -61,21 +63,24 @@ def load_model():
 def main(audio):
     newdf = transform_data(audio)
     loaded_model = load_model()
-    newpred = predict(newdf, loaded_model)
+    newpred, proba = predict(newdf, loaded_model)
     final = get_label(newpred)
     
-    return final
+    return final, {'Probability - Si': proba[1],
+                   'Probability - No': proba[0]}
 
 demo = gr.Interface(
     title = "Autoagent | YES or NO Classification - Layer7",
-    description = "<h3>This model is useful to classify if the user says 'Si' or 'No'. 🎙️ </h3> <br> <b>Record your voice:</b>",
+    description = """<h3>This model is useful to classify if the user says 'Si' or 'No'. 🎙️ </h3>
+                    <img src="https://huggingface.co/spaces/Adrian8as/imagen/resolve/main/output.png" width="350" height="350"/> <br>
+                    <b>Record your voice:</b>""",
     allow_flagging = "never",
     fn = main,
     inputs = gr.Audio(
         sources=["microphone"],
         type="filepath",
     ), 
-    outputs = gr.Textbox(label="Clasification")
+    outputs = [gr.Textbox(label="Clasification"),"label"]
     )
     
 if __name__ == "__main__":