Update app.py

app.py

@@ -13,9 +13,15 @@ labels = learner.dls.vocab
 # Define a function to make predictions
 def predict(text):
     prediction = learner.predict(text)
-    print(prediction)
+    
     pred, pred_idx, probs = prediction[:3]
-    return {labels[i]: float(probs[i]) for i in range(len(labels))}
+
+    negative_label, positive_label = labels
+    
+    return {
+        negative_label: float(probs[0]),
+        positive_label: float(probs[1])
+    }
 
 # Create and launch the Gradio interface
 gr.Interface(