Update app.py

app.py

@@ -3,27 +3,45 @@ from transformers import pipeline
 
 # Pre-load multiple models
 models = {
-    "TinyBERT (NLU)": pipeline("fill-mask", model="prajjwal1/bert-tiny"),
-    "DistilBERT (NLU)": pipeline("fill-mask", model="distilbert-base-uncased"),
-    "ALBERT (NLU)": pipeline("fill-mask", model="albert-base-v2"),
-    "MobileBERT (NLU)": pipeline("fill-mask", model="google/mobilebert-uncased"),
-    "GPT-2 (text gen)": pipeline("text-generation", model="gpt2")
+    "TinyBERT (Fill Mask)": pipeline("fill-mask", model="prajjwal1/bert-tiny"),
+    "DistilBERT (Fill Mask)": pipeline("fill-mask", model="distilbert-base-uncased"),
+    "ALBERT (Fill Mask)": pipeline("fill-mask", model="albert-base-v2"),
+    "MobileBERT (Fill Mask)": pipeline("fill-mask", model="google/mobilebert-uncased"),
+    "GPT-2 (Text Generation)": pipeline("text-generation", model="gpt2")
 }
 
 def run_model(model_name, text):
     pipe = models[model_name]
+
+    # GPT-2 → freeform text generation
     if "GPT-2" in model_name:
-        return pipe(text, max_length=50, do_sample=True)[0]["generated_text"]
+        output = pipe(text, max_length=50, do_sample=True, top_k=50, temperature=0.7)
+        return output[0]["generated_text"]
+
+    # Fill-mask models → require [MASK] token
     else:
-        return str(pipe(text))
+        if "[MASK]" not in text:
+            # If user didn’t include a mask, append one
+            text = text.strip()
+            if not text.endswith("."):
+                text += "."
+            text = text[:-1] + " [MASK]."
+
+        preds = pipe(text, top_k=5)  # top 5 predictions
+        formatted = "\n".join(
+            [f"{p['token_str']} (prob={p['score']:.4f})" for p in preds]
+        )
+        return f"Input: {text}\n\nPredictions:\n{formatted}"
 
 with gr.Blocks() as demo:
-    gr.Markdown("# 🔥 Tiny LLM Playground\nChoose a small model and test it!")
-    
+    gr.Markdown("# 🔥 Tiny LLM Playground\nChoose a small model and test it!\n\n"
+                "💡 For BERT-style models, you can add `[MASK]` in your text, "
+                "or just type normally and I'll add one for you.")
+
     model_choice = gr.Dropdown(list(models.keys()), label="Choose Model")
     text_input = gr.Textbox(label="Enter text or prompt")
-    output = gr.Textbox(label="Output")
-    
+    output = gr.Textbox(label="Output", lines=8)
+
     run_btn = gr.Button("Run")
     run_btn.click(fn=run_model, inputs=[model_choice, text_input], outputs=output)