Update app.py

app.py

@@ -1,15 +1,30 @@
 # 🔍 Masked Word Predictor | CPU-only HF Space
 
 import gradio as gr
-from transformers import pipeline
+from transformers import pipeline, PipelineException
 
 # Load the fill-mask pipeline once at startup
 fill_mask = pipeline("fill-mask", model="distilroberta-base", device=-1)
 
 def predict_mask(sentence: str, top_k: int):
-    if "[MASK]" not in sentence:
-        return [{"sequence": "Error: include [MASK] in your sentence.", "score": 0.0}]
-    preds = fill_mask(sentence, top_k=top_k)
+    # Get the model’s actual mask token (e.g. "<mask>")
+    mask = fill_mask.tokenizer.mask_token
+
+    # Allow users to type [MASK]; convert it under the hood
+    if "[MASK]" in sentence:
+        sentence = sentence.replace("[MASK]", mask)
+
+    # If no mask token present, show error
+    if mask not in sentence:
+        return [{"sequence": f"Error: please include `[MASK]` in your sentence.", "score": 0.0}]
+
+    # Call the pipeline and catch any unexpected exceptions
+    try:
+        preds = fill_mask(sentence, top_k=top_k)
+    except PipelineException as e:
+        return [{"sequence": f"Error: {str(e)}", "score": 0.0}]
+
+    # Format into list-of-dicts for Gradio Dataframe
     return [
         {"sequence": p["sequence"], "score": round(p["score"], 3)}
         for p in preds
@@ -18,20 +33,21 @@ def predict_mask(sentence: str, top_k: int):
 with gr.Blocks(title="🔍 Masked Word Predictor") as demo:
     gr.Markdown(
         "# 🔍 Masked Word Predictor\n"
-        "Enter a sentence with one `[MASK]` token and see the model’s top predictions."
+        "Enter a sentence with one `[MASK]` token and see the top-K model predictions."
     )
 
     with gr.Row():
         sentence = gr.Textbox(
             lines=2,
-            placeholder="The capital of France is [MASK].",
+            placeholder="e.g. The salon’s new color treatment is [MASK].",
             label="Input Sentence"
         )
         top_k = gr.Slider(
-            minimum=1, maximum=10, step=1, value=5,
+            minimum=1, maximum=10, value=5, step=1,
             label="Top K Predictions"
         )
-    predict_btn = gr.Button("Predict", variant="primary")
+
+    predict_btn = gr.Button("Predict 🔍", variant="primary")
 
     results = gr.Dataframe(
         headers=["sequence", "score"],