Update app.py

app.py

@@ -3,7 +3,7 @@ import gradio as gr
 from transformers import pipeline
 
 # ————————————————
-# 1) Load taxonomies
+# Load taxonomies
 # ————————————————
 with open("coarse_labels.json") as f:
     coarse_labels = json.load(f)
@@ -11,128 +11,118 @@ with open("fine_labels.json") as f:
     fine_map = json.load(f)
 
 # ————————————————
-# 2) Available zero-shot models
+# Model choices (5 only)
 # ————————————————
 MODEL_CHOICES = [
     "facebook/bart-large-mnli",
     "roberta-large-mnli",
     "joeddav/xlm-roberta-large-xnli",
-    "mistralai/Mistral-7B-Instruct",
-    "huggyllama/llama-2-7b-chat",
-    "google/flan-t5-large",
-    "google/flan-ul2",
-    "clare-ai/llama-2-13b-instruct",
-    "allenai/longformer-base-4096",
-    "facebook/bart-large-mnli",  # duplicate to test allow_custom_value 
-    "valhalla/t5-base-qa-qg-hl",
-    "EleutherAI/gpt-neox-20b",
-    "EleutherAI/gpt-j-6b",
-    "bigscience/bloom-1b1",
-    "bigscience/bloom-560m",
-    "bigscience/bloom-3b",
-    "Salesforce/codegen-2B-multi",
-    "Salesforce/codegen-350M-multi",
-    "madlag/llama2-7b-finetuned-qa",
-    "tiiuae/falcon-7b-instruct",
-    "tiiuae/falcon-40b-instruct",
-    "milvus/milvus-embed-english",
-    "sentence-transformers/all-MiniLM-L6-v2",
-    "YOUR-OWN-CUSTOM-MODEL"
+    "valhalla/distilbart-mnli-12-4",
+    "educationfoundation/Phantom-7B-JEE"  # placeholder — replace with real phantom model
 ]
 
 # ————————————————
-# Helper: ensure log files exist
+# Ensure log files exist
 # ————————————————
 LOG_FILE      = "logs.csv"
 FEEDBACK_FILE = "feedback.csv"
-for fn, hdr in [(LOG_FILE, ["timestamp","model","question","subject","top3_topics","duration"]),
-                (FEEDBACK_FILE, ["timestamp","question","pred_subject","pred_topics","corrected"])]:
+for fn, hdr in [
+    (LOG_FILE,      ["timestamp","model","question","chosen_subject","top3_topics","duration"]),
+    (FEEDBACK_FILE, ["timestamp","question","subject_feedback","topic_feedback"])
+]:
     if not os.path.exists(fn):
         with open(fn, "w", newline="") as f:
-            writer = csv.writer(f)
-            writer.writerow(hdr)
+            csv.writer(f).writerow(hdr)
 
 # ————————————————
-# 3) Build the interface logic
+# Inference functions
 # ————————————————
-
-def hierarchical_tag(question, model_name):
+def run_stage1(question, model_name):
+    """Return top3 coarse subjects + duration."""
     start = time.time()
-    # 3.1 Instantiate classifier per-run (to change models dynamically)
     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
 
-    # 3.2 Stage 1: coarse label
-    coarse_out = clf(question, candidate_labels=coarse_labels)
-    subject = coarse_out["labels"][0]
-
-    # 3.3 Stage 2: fine labels within chosen subject
+def run_stage2(question, model_name, subject):
+    """Return top3 fine topics + duration."""
+    start = time.time()
+    clf = pipeline("zero-shot-classification", model=model_name)
     fine_labels = fine_map.get(subject, [])
-    fine_out = clf(question, candidate_labels=fine_labels)
-    top3 = fine_out["labels"][:3]
-
-    duration = round(time.time() - start, 3)
-
-    # 3.4 Log the run
+    out = clf(question, candidate_labels=fine_labels)
+    labels, scores = out["labels"][:3], out["scores"][:3]
+    duration = round(time.time()-start,3)
+    # Log combined run
     with open(LOG_FILE, "a", newline="") as f:
         csv.writer(f).writerow([
             time.strftime("%Y-%m-%d %H:%M:%S"),
             model_name,
             question.replace("\n"," "),
             subject,
-            ";".join(top3),
+            ";".join(labels),
             duration
         ])
+    return {lbl: round(score,3) for lbl,score in zip(labels, scores)}, f"⏱ {duration}s"
 
-    # 3.5 Return for display
-    return subject, {lbl: round(score,3) 
-                     for lbl,score in zip(fine_out["labels"][:3],
-                                         fine_out["scores"][:3]
-                                        )}, f"⏱ {duration}s"
-
-def submit_feedback(question, subject, topics, corrected):
-    ts = time.strftime("%Y-%m-%d %H:%M:%S")
+def submit_feedback(question, subject_fb, topic_fb):
     with open(FEEDBACK_FILE, "a", newline="") as f:
-        csv.writer(f).writerow([ts, question.replace("\n"," "), subject, ";".join(topics), corrected])
-    return "Thank you for your feedback!"
+        csv.writer(f).writerow([
+            time.strftime("%Y-%m-%d %H:%M:%S"),
+            question.replace("\n"," "),
+            subject_fb,
+            topic_fb
+        ])
+    return "✅ Feedback recorded!"
 
 # ————————————————
-# 4) Define the Gradio UI
+# Build Gradio UI
 # ————————————————
 with gr.Blocks() as demo:
-    gr.Markdown("## Hierarchical Zero-Shot Tagger with Model Selection & Logging")
+    gr.Markdown("## Hierarchical Zero-Shot Tagger with Subject Toggle & Feedback")
 
     with gr.Row():
         question_input = gr.Textbox(lines=3, label="Enter your question")
         model_input    = gr.Dropdown(
-            label="Choose model",
-            choices=MODEL_CHOICES,
-            value=MODEL_CHOICES[0],
-            allow_custom_value=True
+            choices=MODEL_CHOICES, value=MODEL_CHOICES[0], label="Choose model"
         )
+        go_button      = gr.Button("Run Stage 1")
 
-    run_button = gr.Button("Tag Question")
+    # Stage 1 outputs
+    subj_radio     = gr.Radio(choices=[], label="Top-3 Subjects",
+                              info="Select to re-run Stage 2 for a different subject")
+    stage1_time    = gr.Textbox(label="Stage 1 Time")
 
-    subject_out = gr.Textbox(label="Predicted Subject")
+    go_button.click(
+        fn=lambda q,m: (*run_stage1(q,m),), 
+        inputs=[question_input, model_input],
+        outputs=[subj_radio, stage1_time]
+    )
+
+    # Stage 2 UI
+    go2_button  = gr.Button("Run Stage 2")
     topics_out  = gr.Label(label="Top-3 Topics")
-    time_out    = gr.Textbox(label="Inference Time")
+    stage2_time = gr.Textbox(label="Stage 2 Time")
 
-    run_button.click(
-        hierarchical_tag,
-        inputs=[question_input, model_input],
-        outputs=[subject_out, topics_out, time_out]
+    go2_button.click(
+        fn=run_stage2,
+        inputs=[question_input, model_input, subj_radio],
+        outputs=[topics_out, stage2_time]
     )
 
     gr.Markdown("---")
-    gr.Markdown("### Not quite right? Submit your corrections below:")
+    gr.Markdown("### Feedback / Correction")
 
-    corrected_input = gr.Textbox(lines=1, placeholder="Correct subject;topic1;topic2;topic3")
-    feedback_button = gr.Button("Submit Feedback")
-    feedback_status = gr.Textbox(label="")
+    subject_fb = gr.Textbox(label="Correct Subject")
+    topic_fb   = gr.Textbox(label="Correct Topic(s)")
+    fb_button  = gr.Button("Submit Feedback")
+    fb_status  = gr.Textbox(label="")
 
-    feedback_button.click(
-        submit_feedback,
-        inputs=[question_input, subject_out, topics_out, corrected_input],
-        outputs=[feedback_status]
+    fb_button.click(
+        fn=submit_feedback,
+        inputs=[question_input, subject_fb, topic_fb],
+        outputs=[fb_status]
     )
 
-demo.launch()
+demo.launch(share=True, ssr=False)