Update app.py

app.py

@@ -1,8 +1,10 @@
 import gradio as gr
 from transformers import pipeline
 
-# Pre-load multiple models
-models = {
+# ----------------
+# TEXT MODELS
+# ----------------
+text_models = {
     "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"),
@@ -10,39 +12,77 @@ models = {
     "GPT-2 (Text Generation)": pipeline("text-generation", model="gpt2")
 }
 
-def run_model(model_name, text):
-    pipe = models[model_name]
+def run_text_model(model_name, text):
+    pipe = text_models[model_name]
 
-    # GPT-2 → freeform text generation
     if "GPT-2" in model_name:
         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:
         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
+        preds = pipe(text, top_k=5)
         formatted = "\n".join(
             [f"{p['token_str']} (prob={p['score']:.4f})" for p in preds]
         )
         return f"Input: {text}\n\nPredictions:\n{formatted}"
 
+
+# ----------------
+# IMAGE SEGMENTATION
+# ----------------
+segmentation_pipeline = pipeline(
+    "image-segmentation", model="nvidia/segformer-b0-finetuned-ade-512-512"
+)
+
+def segment_image(image):
+    results = segmentation_pipeline(image)
+    # Gradio AnnotatedImage expects (image, annotations)
+    ann = [(image, r["mask"]) for r in results]
+    return (image, ann)
+
+
+# ----------------
+# SPEECH RECOGNITION
+# ----------------
+asr_pipeline = pipeline("automatic-speech-recognition", model="openai/whisper-tiny")
+
+def transcribe(audio):
+    return asr_pipeline(audio)["text"]
+
+
+# ----------------
+# GRADIO APP
+# ----------------
 with gr.Blocks() as demo:
-    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.")
+    gr.Markdown("# 🔥 Multi-Modal Playground\n"
+                "Try **Tiny LLMs, Image Segmentation, and Speech Models** all in one app!\n\n")
+
+    # TEXT TAB
+    with gr.Tab("Text Models"):
+        model_choice = gr.Dropdown(list(text_models.keys()), label="Choose Model")
+        text_input = gr.Textbox(label="Enter text or prompt")
+        text_output = gr.Textbox(label="Output", lines=8)
+        run_btn = gr.Button("Run")
+        run_btn.click(fn=run_text_model, inputs=[model_choice, text_input], outputs=text_output)
 
-    model_choice = gr.Dropdown(list(models.keys()), label="Choose Model")
-    text_input = gr.Textbox(label="Enter text or prompt")
-    output = gr.Textbox(label="Output", lines=8)
+    # IMAGE TAB
+    with gr.Tab("Image Segmentation"):
+        img_in = gr.Image(type="pil", label="Upload an Image")
+        img_out = gr.AnnotatedImage(label="Segmented Output")
+        seg_btn = gr.Button("Segment Objects")
+        seg_btn.click(fn=segment_image, inputs=img_in, outputs=img_out)
 
-    run_btn = gr.Button("Run")
-    run_btn.click(fn=run_model, inputs=[model_choice, text_input], outputs=output)
+    # AUDIO TAB
+    with gr.Tab("Speech Recognition"):
+        audio_in = gr.Audio(sources=["microphone", "upload"], type="filepath", label="Upload or record audio")
+        audio_out = gr.Textbox(label="Transcription")
+        asr_btn = gr.Button("Transcribe")
+        asr_btn.click(fn=transcribe, inputs=audio_in, outputs=audio_out)
 
 demo.launch()