Update app.py

app.py

@@ -1,51 +1,94 @@
 import gradio as gr
 from transformers import pipeline
 
-# Load multiple open-source models (lightweight for CPU-basic)
+# ===============================
+# Load open-source text generation models
+# ===============================
 models = {
-    "DistilGPT-2": pipeline("text-generation", model="distilgpt2"),
-    "GPT2 (Small)": pipeline("text-generation", model="gpt2"),
-    "DialoGPT-small": pipeline("text-generation", model="microsoft/DialoGPT-small"),
-    "OPT-350M": pipeline("text-generation", model="facebook/opt-350m"),
-    "Bloom-560M": pipeline("text-generation", model="bigscience/bloom-560m"),
+    "DistilGPT-2": "distilgpt2",
+    "GPT2 (Small)": "gpt2",
+    "DialoGPT-small": "microsoft/DialoGPT-small",
+    "OPT-350M": "facebook/opt-350m",
+    "Bloom-560M": "bigscience/bloom-560m",
+    "GPT-Neo-125M": "EleutherAI/gpt-neo-125M",
+    "Falcon-RW-1B": "tiiuae/falcon-rw-1b",
+    "Flan-T5-Small": "google/flan-t5-small",
+    "Flan-T5-Base": "google/flan-t5-base",
+    "Phi-2": "microsoft/phi-2"
 }
 
+generators = {name: pipeline("text-generation", model=mdl) 
+              if "flan" not in mdl.lower() and "bart" not in mdl.lower()
+              else pipeline("text2text-generation", model=mdl)
+              for name, mdl in models.items()}
+
+# Summarizer model
+summarizer = pipeline("text2text-generation", model="google/flan-t5-base")
+
+# ===============================
+# Function to query all models
+# ===============================
 def compare_models(user_input, max_new_tokens=100, temperature=0.7):
-    results = {}
-    for name, generator in models.items():
+    raw_outputs = {}
+    clean_outputs = {}
+    for name, generator in generators.items():
         try:
-            output = generator(user_input, max_new_tokens=max_new_tokens, temperature=temperature)[0]["generated_text"]
-            results[name] = output
+            if "text-generation" in generator.task:
+                output = generator(
+                    user_input,
+                    max_new_tokens=max_new_tokens,
+                    temperature=temperature
+                )[0]["generated_text"]
+            else:  # Flan models etc
+                output = generator(user_input, max_new_tokens=max_new_tokens)[0]["generated_text"]
+
+            raw_outputs[name] = output
+
+            # Summarize the answer to improve clarity
+            summary = summarizer("Summarize this: " + output, max_new_tokens=60)[0]["generated_text"]
+            clean_outputs[name] = summary
+
         except Exception as e:
-            results[name] = f"⚠️ Error: {str(e)}"
-    return [results[m] for m in models.keys()]
+            raw_outputs[name] = f"⚠️ Error: {str(e)}"
+            clean_outputs[name] = "N/A"
+
+    return [raw_outputs[m] for m in models.keys()], [clean_outputs[m] for m in models.keys()]
 
+# ===============================
+# Gradio UI
+# ===============================
 with gr.Blocks(css="style.css") as demo:
     gr.Markdown("## 🤖 Open-Source Model Comparator\n"
-                "Type a prompt once, see how different open-source LLMs respond side-by-side.")
+                "Compare outputs from multiple open-source LLMs side by side.\n"
+                "Includes a raw output and a cleaned summary (via Flan-T5).")
 
     with gr.Row():
-        user_input = gr.Textbox(label="Your prompt", placeholder="Ask something like 'Write a short poem about the stars'...", lines=2)
+        user_input = gr.Textbox(label="Your prompt", placeholder="Try: 'Explain quantum computing in simple terms'", lines=2)
         generate_btn = gr.Button("Generate", variant="primary")
 
     with gr.Row():
         max_tokens = gr.Slider(20, 200, value=100, step=10, label="Max new tokens")
         temp = gr.Slider(0.1, 1.0, value=0.7, step=0.1, label="Creativity (temperature)")
 
+    gr.Markdown("### 🔎 Raw Outputs")
+    with gr.Row():
+        raw_boxes = [gr.Textbox(label=name, elem_classes="output-box", interactive=False) for name in models.keys()]
+
+    gr.Markdown("### ✨ Cleaned Summaries (Flan-T5)")
     with gr.Row():
-        outputs = [gr.Textbox(label=name, elem_classes="output-box", interactive=False) for name in models.keys()]
+        clean_boxes = [gr.Textbox(label=f"{name} (Summary)", elem_classes="output-box", interactive=False) for name in models.keys()]
 
     examples = [
         ["Explain quantum computing in simple terms."],
         ["Write a haiku about autumn leaves."],
         ["What are the pros and cons of nuclear energy?"],
         ["Describe a futuristic city in the year 2200."],
-        ["Write a funny short story about a robot learning to cook."],
+        ["Write a funny short story about a robot learning to cook."]
     ]
     gr.Examples(examples=examples, inputs=[user_input])
 
-    generate_btn.click(compare_models, inputs=[user_input, max_tokens, temp], outputs=outputs)
-    user_input.submit(compare_models, inputs=[user_input, max_tokens, temp], outputs=outputs)
+    generate_btn.click(compare_models, inputs=[user_input, max_tokens, temp], outputs=raw_boxes + clean_boxes)
+    user_input.submit(compare_models, inputs=[user_input, max_tokens, temp], outputs=raw_boxes + clean_boxes)
 
 if __name__ == "__main__":
     demo.launch()