Update app.py

app.py

@@ -38,13 +38,18 @@ for fn, hdr in [
 # Inference functions
 # ————————————————
 def run_stage1(question, model_name):
-    """Return top3 coarse subjects + duration."""
     start = time.time()
     clf = pipeline("zero-shot-classification", model=model_name)
     out = clf(question, candidate_labels=coarse_labels)
     labels, scores = out["labels"][:3], out["scores"][:3]
-    duration = round(time.time()-start,3)
-    return labels, duration
+    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
+    radio_update = gr.update(choices=labels, value=labels[0])
+
+    return label_dict, radio_update, f"⏱ {duration}s"
 
 def run_stage2(question, model_name, subject):
     """Return top3 fine topics + duration."""
@@ -84,20 +89,17 @@ with gr.Blocks() as demo:
 
     with gr.Row():
         question_input = gr.Textbox(lines=3, label="Enter your question")
-        model_input    = gr.Dropdown(
-            choices=MODEL_CHOICES, value=MODEL_CHOICES[0], label="Choose model"
-        )
+        model_input    = gr.Dropdown(choices=MODEL_CHOICES, value=MODEL_CHOICES[0], label="Choose model")
         go_button      = gr.Button("Run Stage 1")
 
-    # Stage 1 outputs
-    subj_radio     = gr.Radio(choices=[], label="Top-3 Subjects",
-                              info="Select to re-run Stage 2 for a different subject")
+    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=lambda q,m: (*run_stage1(q,m),), 
+        fn=run_stage1,
         inputs=[question_input, model_input],
-        outputs=[subj_radio, stage1_time]
+        outputs=[subject_out, subj_radio, stage1_time]
     )
 
     # Stage 2 UI