Update app.py

app.py

@@ -2,9 +2,25 @@ import gradio as gr
 from transformers import pipeline
 
 def mask(text):
-  mask_model = pipeline("fill-mask", model="google-bert/bert-base-uncased")
-  output = mask_model(text)
-  return output[0]['sequence']
+    mask_model = pipeline("fill-mask", model="google-bert/bert-base-uncased")
+
+    # Extract labels (classes) and scores from predictions
+    labels = [result["label"] for result in mask_model]
+    scores = [result["score"] for result in mask_model]
+
+    # Find the index of the best prediction (highest score)
+    best_prediction_idx = scores.index(max(scores))
+
+    # Create a dictionary with results
+    response_dict = {
+        "Original Text": text,
+        "All Predictions": labels,
+        "Best Prediction": f"**{labels[best_prediction_idx]}** (Score: {scores[best_prediction_idx]:.4f})"
+    }
+
+    return response_dict
+    # output = mask_model(text)
+    # return output[0]['sequence']
 
 # Gradio UI
 examples=[['Today I went to [MASK] after I got out of bed.']]
@@ -17,4 +33,4 @@ ui = gr.Interface(
     description="Enter some text with [MASK] to let BERT guess the missing word!",
     examples=examples
     )
-ui.launch(debug=True)
+ui.launch(debug=True)