Update app.py

app.py

@@ -3,24 +3,9 @@ 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']
 
-    # 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.']]