Update app.py

app.py

@@ -2,101 +2,65 @@ import gradio as gr
 from transformers import pipeline
 
 # ===============================
-# Model dictionary (lazy loaded)
+# Load only text-generation models (simpler, stable)
 # ===============================
-model_names = {
-    "DistilGPT-2": "distilgpt2",
-    "Bloom-560M": "bigscience/bloom-560m",
-    "OPT-350M": "facebook/opt-350m",
-    "Flan-T5-Base": "google/flan-t5-base",
-    "Phi-2": "microsoft/phi-2"
+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"),
+    "GPT-Neo-125M": pipeline("text-generation", model="EleutherAI/gpt-neo-125M"),
+    "Falcon-RW-1B": pipeline("text-generation", model="tiiuae/falcon-rw-1b"),
 }
 
-loaded_models = {}
-summarizer = None  # Flan-T5 for cleanup
-
-# ===============================
-# Lazy-load helper
-# ===============================
-def get_model(name):
-    if name not in loaded_models:
-        mdl = model_names[name]
-        if "flan" in mdl.lower():
-            loaded_models[name] = pipeline("text2text-generation", model=mdl)
-        else:
-            loaded_models[name] = pipeline("text-generation", model=mdl)
-    return loaded_models[name]
-
-def get_summarizer():
-    global summarizer
-    if summarizer is None:
-        summarizer = pipeline("text2text-generation", model="google/flan-t5-base")
-    return summarizer
-
-# ===============================
-# Compare function
-# ===============================
-def compare_models(user_input, max_new_tokens=100, temperature=0.7):
-    raw_outputs, clean_outputs = {}, {}
-    for name in model_names.keys():
+def compare_models(user_input, max_new_tokens=100, temperature=0.7, top_p=0.95):
+    results = {}
+    for name, generator in models.items():
         try:
-            generator = get_model(name)
-            if generator.task == "text-generation":
-                output = generator(
-                    user_input,
-                    max_new_tokens=max_new_tokens,
-                    temperature=temperature
-                )[0]["generated_text"]
-            else:  # text2text-generation (Flan)
-                output = generator(user_input, max_new_tokens=max_new_tokens)[0]["generated_text"]
-
-            raw_outputs[name] = output
-
-            # Summarize
-            summary = get_summarizer()("Summarize this: " + output, max_new_tokens=60)[0]["generated_text"]
-            clean_outputs[name] = summary
-
+            output = generator(
+                user_input,
+                max_new_tokens=max_new_tokens,
+                temperature=temperature,
+                top_p=top_p,
+                do_sample=True
+            )[0]["generated_text"]
+            results[name] = output
         except Exception as e:
-            raw_outputs[name] = f"⚠️ Error: {str(e)}"
-            clean_outputs[name] = "N/A"
-
-    return [raw_outputs[m] for m in model_names.keys()], [clean_outputs[m] for m in model_names.keys()]
+            results[name] = f"⚠️ Error: {str(e)}"
+    return [results[m] for m in models.keys()]
 
 # ===============================
 # Gradio UI
 # ===============================
 with gr.Blocks(css="style.css") as demo:
     gr.Markdown("## 🤖 Open-Source Model Comparator\n"
-                "Compare outputs from open-source LLMs side by side.\n"
-                "Raw output + a cleaned summary from Flan-T5.")
+                "Compare outputs from multiple open-source LLMs side by side.\n"
+                "These are raw, unfiltered outputs from Hugging Face models.")
 
     with gr.Row():
-        user_input = gr.Textbox(label="Your prompt", placeholder="Try: 'Explain quantum computing in simple terms'", lines=2)
+        user_input = gr.Textbox(label="Your prompt", placeholder="Ask something like 'Write a short poem about the stars'...", 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)")
+        topp = gr.Slider(0.5, 1.0, value=0.95, step=0.05, label="Nucleus sampling (top_p)")
 
-    gr.Markdown("### 🔎 Raw Outputs")
-    with gr.Row():
-        raw_boxes = [gr.Textbox(label=name, elem_classes="output-box", interactive=False) for name in model_names.keys()]
-
-    gr.Markdown("### ✨ Cleaned Summaries (Flan-T5)")
     with gr.Row():
-        clean_boxes = [gr.Textbox(label=f"{name} (Summary)", elem_classes="output-box", interactive=False) for name in model_names.keys()]
+        outputs = [gr.Textbox(label=name, 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=raw_boxes + clean_boxes)
-    user_input.submit(compare_models, inputs=[user_input, max_tokens, temp], outputs=raw_boxes + clean_boxes)
+    generate_btn.click(compare_models, inputs=[user_input, max_tokens, temp, topp], outputs=outputs)
+    user_input.submit(compare_models, inputs=[user_input, max_tokens, temp, topp], outputs=outputs)
 
 if __name__ == "__main__":
     demo.launch()