Update app.py

app.py

@@ -21,11 +21,13 @@ sp = load_sp_model()
 # ---------------------- Prediction Function ----------------------
 @torch.no_grad()
 def predict_next_words(text, max_predictions=3):
-    if not text.strip():
+    """Predict up to max_predictions next words."""
+    text = text.strip()
+    if not text:
         return []
 
-    token_ids = sp.encode(text.strip(), out_type=int)
-    if len(token_ids) == 0:
+    token_ids = sp.encode(text, out_type=int)
+    if not token_ids:
         return []
 
     input_seq = torch.tensor(token_ids, dtype=torch.long).unsqueeze(0).to(device)
@@ -38,13 +40,23 @@ def predict_next_words(text, max_predictions=3):
 
 # ---------------------- Gradio App Functions ----------------------
 def submit_and_predict(text):
+    # Get predictions and ensure exactly 3 by padding empty strings.
     suggestions = predict_next_words(text)
-    suggestions += [""] * (3 - len(suggestions))  # Ensure 3 buttons always
-    return suggestions
+    suggestions += [""] * (3 - len(suggestions))
+    
+    # Return an array of Gradio "update" objects so we can make them visible or hidden.
+    updates = []
+    for s in suggestions:
+        if s:  # Valid prediction
+            updates.append(gr.update(value=s, visible=True))
+        else:  # No prediction
+            updates.append(gr.update(value="", visible=False))
+    return updates
 
 def append_suggestion(text, suggestion):
+    # Only append if not empty.
     if suggestion:
-        text += suggestion + " "
+        text = text.rstrip() + " " + suggestion + " "
     return text
 
 # ---------------------- Gradio Interface ----------------------
@@ -58,12 +70,14 @@ with gr.Blocks(title="Next Word Predictor") as app:
     with gr.Row():
         suggestion_buttons = [gr.Button(visible=False) for _ in range(3)]
 
+    # 1. When user clicks 'Submit', run submit_and_predict to get suggestions.
     submit_btn.click(
         fn=submit_and_predict,
         inputs=text_input,
         outputs=suggestion_buttons,
     )
 
+    # 2. Each suggestion button appends the chosen word to the main text.
     for btn in suggestion_buttons:
         btn.click(
             fn=append_suggestion,