Update app.py

app.py

@@ -2,39 +2,35 @@
 
 import gradio as gr
 from transformers import pipeline
-from transformers.pipelines.base import PipelineException  # correct import
+from transformers.pipelines.base import PipelineException
 
-# Load the fill-mask pipeline once at startup
+# 1. Load fill-mask pipeline once
 fill_mask = pipeline("fill-mask", model="distilroberta-base", device=-1)
 
 def predict_mask(sentence: str, top_k: int):
-    # Get the model’s actual mask token (e.g. "<mask>")
+    # 2. Get the 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)
+    # 3. Allow users to type [MASK]
+    sentence = sentence.replace("[MASK]", mask)
 
-    # If no mask token present, show error
+    # 4. Validate presence of mask
     if mask not in sentence:
-        return [{"sequence": "Error: please include `[MASK]` in your sentence.", "score": 0.0}]
+        return [["Error: please include `[MASK]` in your sentence.", 0.0]]
 
-    # Call the pipeline and catch any pipeline-specific exceptions
+    # 5. Run the pipeline safely
     try:
         preds = fill_mask(sentence, top_k=top_k)
     except PipelineException as e:
-        return [{"sequence": f"Error: {str(e)}", "score": 0.0}]
+        return [[f"Error: {str(e)}", 0.0]]
 
-    # Format into list-of-dicts for Gradio Dataframe
-    return [
-        {"sequence": p["sequence"], "score": round(p["score"], 3)}
-        for p in preds
-    ]
+    # 6. Return list-of-lists for the Table
+    return [[p["sequence"], round(p["score"], 3)] for p in preds]
 
 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 top-K model predictions."
+        "Enter a sentence with one `[MASK]` token and see the top-K completions."
     )
 
     with gr.Row():
@@ -44,24 +40,22 @@ with gr.Blocks(title="🔍 Masked Word Predictor") as demo:
             label="Input Sentence"
         )
         top_k = gr.Slider(
-            minimum=1, maximum=10, value=5, step=1,
+            minimum=1, maximum=10, step=1, value=5,
             label="Top K Predictions"
         )
 
     predict_btn = gr.Button("Predict 🔍", variant="primary")
 
-    results = gr.Dataframe(
-        headers=["sequence", "score"],
-        datatype=["str", "number"],
-        wrap=True,
-        interactive=False,
+    # ← Use Table here (list-of-lists) instead of Dataframe
+    results_table = gr.Table(
+        headers=["Sequence", "Score"],
         label="Predictions"
     )
 
     predict_btn.click(
-        predict_mask,
+        fn=predict_mask,
         inputs=[sentence, top_k],
-        outputs=results
+        outputs=results_table
     )
 
 if __name__ == "__main__":