Update app.py

app.py

@@ -2,65 +2,73 @@ import gradio as gr
 from transformers import pipeline
 
 # ===============================
-# Load open-source text generation models
+# Model dictionary (lazy loaded)
 # ===============================
-models = {
+model_names = {
     "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",
+    "OPT-350M": "facebook/opt-350m",
     "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()}
+loaded_models = {}
+summarizer = None  # Flan-T5 for cleanup
 
-# Summarizer model
-summarizer = pipeline("text2text-generation", model="google/flan-t5-base")
+# ===============================
+# 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
 
 # ===============================
-# Function to query all models
+# Compare function
 # ===============================
 def compare_models(user_input, max_new_tokens=100, temperature=0.7):
-    raw_outputs = {}
-    clean_outputs = {}
-    for name, generator in generators.items():
+    raw_outputs, clean_outputs = {}, {}
+    for name in model_names.keys():
         try:
-            if "text-generation" in generator.task:
+            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:  # Flan models etc
+            else:  # text2text-generation (Flan)
                 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"]
+            # Summarize
+            summary = get_summarizer()("Summarize this: " + output, max_new_tokens=60)[0]["generated_text"]
             clean_outputs[name] = summary
 
         except Exception as e:
             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()]
+    return [raw_outputs[m] for m in model_names.keys()], [clean_outputs[m] for m in model_names.keys()]
 
 # ===============================
 # Gradio UI
 # ===============================
 with gr.Blocks(css="style.css") as demo:
     gr.Markdown("## 🤖 Open-Source Model Comparator\n"
-                "Compare outputs from multiple open-source LLMs side by side.\n"
-                "Includes a raw output and a cleaned summary (via Flan-T5).")
+                "Compare outputs from open-source LLMs side by side.\n"
+                "Raw output + a cleaned summary from Flan-T5.")
 
     with gr.Row():
         user_input = gr.Textbox(label="Your prompt", placeholder="Try: 'Explain quantum computing in simple terms'", lines=2)
@@ -72,11 +80,11 @@ with gr.Blocks(css="style.css") as demo:
 
     gr.Markdown("### 🔎 Raw Outputs")
     with gr.Row():
-        raw_boxes = [gr.Textbox(label=name, elem_classes="output-box", interactive=False) for name in models.keys()]
+        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 models.keys()]
+        clean_boxes = [gr.Textbox(label=f"{name} (Summary)", elem_classes="output-box", interactive=False) for name in model_names.keys()]
 
     examples = [
         ["Explain quantum computing in simple terms."],