Update app.py

app.py

@@ -2,6 +2,11 @@ import json, time, csv, os
 import gradio as gr
 from transformers import pipeline
 
+PIPELINES = {
+    name: pipeline("zero-shot-classification", model=name)
+    for name in MODEL_CHOICES
+}
+
 # ————————————————
 # Load taxonomies
 # ————————————————
@@ -38,28 +43,38 @@ for fn, hdr in [
 # Inference functions
 # ————————————————
 def run_stage1(question, model_name):
+    if not question or not question.strip():
+        return {}, gr.update(choices=[]), ""
     start = time.time()
-    clf = pipeline("zero-shot-classification", model=model_name)
+    clf = PIPELINES[model_name]
     out = clf(question, candidate_labels=coarse_labels)
     labels, scores = out["labels"][:3], out["scores"][:3]
     duration = round(time.time() - start, 3)
 
-    # Build the label mapping
-    label_dict = {lbl: round(score,3) for lbl,score in zip(labels, scores)}
-    # Prepare Radio update
+    # Prepare outputs
+    subject_dict = {lbl: round(score,3) for lbl,score in zip(labels, scores)}
     radio_update = gr.update(choices=labels, value=labels[0])
+    time_str    = f"⏱ {duration}s"
+
+    return subject_dict, radio_update, time_str
 
-    return label_dict, radio_update, f"⏱ {duration}s"
 
 def run_stage2(question, model_name, subject):
-    """Return top3 fine topics + duration."""
-    start = time.time()
-    clf = pipeline("zero-shot-classification", model=model_name)
+    # 1) Validate inputs
+    if not question or not question.strip():
+        return {}, "No question provided", ""
     fine_labels = fine_map.get(subject, [])
+    if not fine_labels:
+        return {}, f"No topics found for '{subject}'", ""
+
+    # 2) Inference (fast, using preloaded pipeline)
+    start = time.time()
+    clf = PIPELINES[model_name]
     out = clf(question, candidate_labels=fine_labels)
     labels, scores = out["labels"][:3], out["scores"][:3]
-    duration = round(time.time()-start,3)
-    # Log combined run
+    duration = round(time.time() - start, 3)
+
+    # 3) Logging
     with open(LOG_FILE, "a", newline="") as f:
         csv.writer(f).writerow([
             time.strftime("%Y-%m-%d %H:%M:%S"),
@@ -69,7 +84,11 @@ def run_stage2(question, model_name, subject):
             ";".join(labels),
             duration
         ])
-    return {lbl: round(score,3) for lbl,score in zip(labels, scores)}, f"⏱ {duration}s"
+
+    # 4) Return topics + time
+    topic_dict = {lbl: round(score,3) for lbl,score in zip(labels, scores)}
+    return topic_dict, f"⏱ {duration}s"
+
 
 def submit_feedback(question, subject_fb, topic_fb):
     with open(FEEDBACK_FILE, "a", newline="") as f:
@@ -92,10 +111,10 @@ with gr.Blocks() as demo:
         model_input    = gr.Dropdown(choices=MODEL_CHOICES, value=MODEL_CHOICES[0], label="Choose model")
         go_button      = gr.Button("Run Stage 1")
 
-    subject_out    = gr.Label(num_top_classes=3, label="Top-3 Subjects")
-    subj_radio     = gr.Radio(choices=[], label="Select Subject for Stage 2")
-    stage1_time    = gr.Textbox(label="Stage 1 Time")
-
+    subject_out = gr.Label(num_top_classes=3, label="Top-3 Subjects")
+    subj_radio  = gr.Radio(choices=[], label="Select Subject for Stage 2")
+    stage1_time = gr.Textbox(label="Stage 1 Time")
+    
     go_button.click(
         fn=run_stage1,
         inputs=[question_input, model_input],