Update app.py

app.py

@@ -1,8 +1,13 @@
 import json
 from functools import partial
+
 import gradio as gr
 from transformers import pipeline
 
+
+# -----------------------------
+# Load module repository
+# -----------------------------
 with open("modules.json", "r", encoding="utf-8") as f:
     MODULES = json.load(f)["modules"]
 
@@ -10,14 +15,46 @@ GENERATORS = [m for m in MODULES if m["type"] == "generator"]
 CHECKERS = {m["id"]: m for m in MODULES if m["type"] == "checker"}
 GEN_BY_ID = {m["id"]: m for m in GENERATORS}
 
+# Friendly names for classifier
+MODULE_LABELS = {
+    "Analysis Note": "analysis_note_v1",
+    "Document Explainer": "document_explainer_v1",
+    "Strategy Memo": "strategy_memo_v1",
+    "Message / Post Reply": "message_reply_v1",
+    "Profile / Application": "profile_application_v1",
+    "System / Architecture Blueprint": "system_blueprint_v1",
+    "Modular Brainstorm": "modular_brainstorm_v1",
+}
+
+LABEL_LIST = list(MODULE_LABELS.keys())
+
+
+# -----------------------------
+# Models
+# -----------------------------
+# 1) Generator engine (same as before)
 llm = pipeline("text-generation", model="gpt2", max_new_tokens=512)
 
-def call_llm(prompt):
-    o = llm(prompt, max_new_tokens=512, do_sample=False)[0]["generated_text"]
-    return o[len(prompt):].strip() if o.startswith(prompt) else o
+# 2) Task classifier (zero-shot)
+#    You can swap this to another classification model later.
+classifier = pipeline(
+    "zero-shot-classification",
+    model="facebook/bart-large-mnli"
+)
+
 
-def generator_prompt(mid, *inputs):
-    m = GEN_BY_ID[mid]
+# -----------------------------
+# Low-level helpers
+# -----------------------------
+def call_llm(prompt: str) -> str:
+    """Call the generator LLM and strip prompt prefix if echoed."""
+    out = llm(prompt, max_new_tokens=512, do_sample=False)[0]["generated_text"]
+    return out[len(prompt):].strip() if out.startswith(prompt) else out
+
+
+def generator_prompt(module_id: str, *inputs: str) -> str:
+    """Build a structured prompt for a generator module."""
+    m = GEN_BY_ID[module_id]
     keys = list(m["input_placeholders"].keys())
     vals = {k: inputs[i] if i < len(inputs) else "" for i, k in enumerate(keys)}
     secs = m["output_sections"]
@@ -38,11 +75,15 @@ def generator_prompt(mid, *inputs):
         p.append("")
     return "\n".join(p)
 
-def checker_prompt(cid, *vals):
-    c = CHECKERS[cid]
+
+def checker_prompt(checker_id: str, *vals: str) -> str:
+    """Build a prompt for a checker module."""
+    c = CHECKERS[checker_id]
     secs = c["output_sections"]
+
     if len(vals) < 2:
-        orig, draft = "", vals[0] if vals else ""
+        orig = ""
+        draft = vals[0] if vals else ""
     else:
         orig = "\n\n".join(vals[:-1])
         draft = vals[-1]
@@ -65,16 +106,96 @@ def checker_prompt(cid, *vals):
         p.append("")
     return "\n".join(p)
 
-def run_generator(mid, *inputs):
+
+def run_generator(mid: str, *inputs: str) -> str:
     return call_llm(generator_prompt(mid, *inputs))
 
-def run_checker(cid, *inputs):
+
+def run_checker(cid: str, *inputs: str) -> str:
     return call_llm(checker_prompt(cid, *inputs))
 
+
+# -----------------------------
+# Task classifier
+# -----------------------------
+def classify_task(task_text: str):
+    """
+    Classify a free-form task description into one of the Modular Intelligence modules.
+
+    Returns:
+        predicted_label: human-friendly label (e.g. "Strategy Memo")
+        module_id: internal id (e.g. "strategy_memo_v1")
+        scores_text: formatted confidence scores per label
+    """
+    if not task_text.strip():
+        return "No input", "", ""
+
+    res = classifier(
+        task_text,
+        candidate_labels=LABEL_LIST,
+        multi_label=False
+    )
+    # Zero-shot pipeline returns labels sorted by score descending
+    predicted_label = res["labels"][0]
+    module_id = MODULE_LABELS[predicted_label]
+
+    # Build a small score table
+    lines = []
+    for label, score in zip(res["labels"], res["scores"]):
+        lines.append(f"{label}: {score:.2f}")
+    scores_text = "\n".join(lines)
+
+    return predicted_label, module_id, scores_text
+
+
+# -----------------------------
+# UI
+# -----------------------------
 def build_ui():
     with gr.Blocks(title="Modular Intelligence") as demo:
-        gr.Markdown("# Modular Intelligence Demo")
+        gr.Markdown(
+            "# Modular Intelligence Demo\n"
+            "Choose a module directly, or use the **Auto-route** tab to classify your task."
+        )
+
+        # ---- Auto-route tab (task classifier) ----
+        with gr.Tab("Auto-route (Task Classifier)"):
+            gr.Markdown(
+                "Paste any task or question. The classifier will suggest "
+                "which Modular Intelligence module to use."
+            )
+            task_box = gr.Textbox(
+                label="Describe what you want to do",
+                placeholder="E.g. 'I want a structured memo on options for exiting a JV under time pressure'...",
+                lines=6
+            )
+            predicted_label_box = gr.Textbox(
+                label="Predicted module (human-readable)",
+                interactive=False
+            )
+            predicted_id_box = gr.Textbox(
+                label="Internal module id",
+                interactive=False
+            )
+            scores_box = gr.Textbox(
+                label="Classifier scores",
+                interactive=False,
+                lines=10
+            )
+
+            classify_button = gr.Button("Classify task")
+            classify_button.click(
+                fn=classify_task,
+                inputs=[task_box],
+                outputs=[predicted_label_box, predicted_id_box, scores_box]
+            )
 
+            gr.Markdown(
+                "You can then go to the corresponding module tab below and "
+                "fill in its inputs using this guidance."
+            )
+
+        # ---- One tab per generator module ----
         for m in GENERATORS:
             with gr.Tab(m["label"]):
                 gr.Markdown(m["description"])
@@ -100,8 +221,10 @@ def build_ui():
                             inputs=inputs + [out],
                             outputs=chk
                         )
+
     return demo
 
+
 if __name__ == "__main__":
     app = build_ui()
     app.launch()