Update app.py

app.py

@@ -1,59 +1,168 @@
 import json
+import os
+import torch
 from functools import partial
 import gradio as gr
-from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
-
-# -------------------------------------------------------------
-# Load Modules
-# -------------------------------------------------------------
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    pipeline
+)
+
+# =============================================================
+# LOAD MODULES.JSON
+# =============================================================
 with open("modules.json", "r", encoding="utf-8") as f:
     MODULES = json.load(f)["modules"]
 
 GENERATORS = [m for m in MODULES if m.get("type") == "generator"]
 CHECKERS = {m["id"]: m for m in MODULES if m.get("type") == "checker"}
 GEN_BY_ID = {m["id"]: m for m in GENERATORS}
+
 LABEL_TO_ID = {m["label"]: m["id"] for m in GENERATORS}
 LABEL_LIST = list(LABEL_TO_ID.keys())
 
-# -------------------------------------------------------------
-# Load Model (base LLM)
-# Swappable engine: GPT-2 / Llama / Mistral etc.
-# -------------------------------------------------------------
+
+# =============================================================
+# BASE MODEL (ENGINE) — Can be swapped
+# =============================================================
 tokenizer = AutoTokenizer.from_pretrained("openai-community/gpt2")
 model = AutoModelForCausalLM.from_pretrained("openai-community/gpt2")
 llm = pipeline("text-generation", model=model, tokenizer=tokenizer, max_new_tokens=300)
 
-# -------------------------------------------------------------
-# Automatic Router Components
-# -------------------------------------------------------------
-from router.rules import rule_router
-from router.zero_shot import classify_task  # zero-shot classifier
-
-# -------------------------------------------------------------
-# Domain Adapters (LoRA)
-# -------------------------------------------------------------
-from domain_heads.loader import load_adapter  # load domain-specific adapter
 
-# -------------------------------------------------------------
-# Reasoning Scaffolds
-# -------------------------------------------------------------
-from reasoning_scaffolds.cot import apply_cot
-from reasoning_scaffolds.critique_loop import critique_and_refine
-
-
-# -------------------------------------------------------------
-# Helper: LLM call
-# -------------------------------------------------------------
+# =============================================================
+# HYBRID ROUTER (RULES + ZERO-SHOT)
+# =============================================================
+
+# ----------- RULE-BASED ROUTER -----------
+RULES = [
+    ("contract", "document_explainer_v1"),
+    ("agreement", "document_explainer_v1"),
+    ("policy", "document_explainer_v1"),
+    ("judgment", "document_explainer_v1"),
+    ("options", "strategy_memo_v1"),
+    ("trade-off", "strategy_memo_v1"),
+    ("recommendation", "strategy_memo_v1"),
+    ("compare", "strategy_memo_v1"),
+    ("system", "system_blueprint_v1"),
+    ("architecture", "system_blueprint_v1"),
+    ("flow", "system_blueprint_v1"),
+    ("analysis", "analysis_note_v1"),
+    ("summarize", "analysis_note_v1"),
+    ("explain", "analysis_note_v1"),
+]
+
+def rule_router(text: str):
+    t = text.lower()
+    for keyword, module_id in RULES:
+        if keyword in t:
+            return module_id
+    return None
+
+
+# ----------- ZERO-SHOT ROUTER -----------
+zero_shot_classifier = pipeline(
+    "zero-shot-classification",
+    model="facebook/bart-large-mnli"
+)
+
+def zero_shot_route(text):
+    res = zero_shot_classifier(text, candidate_labels=LABEL_LIST, multi_label=False)
+    label = res["labels"][0]
+    module_id = LABEL_TO_ID[label]
+    scores = "\n".join([f"{l}: {s:.2f}" for l, s in zip(res["labels"], res["scores"])])
+    return label, module_id, scores
+
+
+# ----------- HYBRID ROUTE CALL -----------
+def hybrid_route(task: str):
+    if not task.strip():
+        return "No input", "", ""
+
+    route = rule_router(task)
+    if route:
+        return GEN_BY_ID[route]["label"], route, "Rule-based match"
+
+    return zero_shot_route(task)
+
+
+# =============================================================
+# DOMAIN HEAD LOADER (LoRA-STYLE ADAPTERS)
+# =============================================================
+ADAPTER_PATHS = {
+    "legal": "domain_heads/legal_head.pt",
+    "strategy": "domain_heads/strategy_head.pt",
+    "analysis": "domain_heads/analysis_head.pt",
+    "systems": "domain_heads/systems_head.pt",
+}
+
+def load_domain_adapter(domain: str):
+    if domain not in ADAPTER_PATHS:
+        return
+
+    path = ADAPTER_PATHS[domain]
+    if not os.path.exists(path):
+        return
+
+    adapter = torch.load(path, map_location="cpu")
+    with torch.no_grad():
+        for name, param in model.named_parameters():
+            if name in adapter:
+                param += adapter[name]
+
+
+# =============================================================
+# REASONING SCAFFOLDS
+# =============================================================
+
+# ----------- CHAIN-OF-THOUGHT -----------
+def apply_cot(prompt: str) -> str:
+    return (
+        "Think step-by-step. Explain your reasoning before answering.\n\n"
+        + prompt
+        + "\n\nNow think step-by-step and answer:"
+    )
+
+# ----------- CRITIQUE + REFINE LOOP -----------
+critic = pipeline(
+    "text-generation",
+    model="openai-community/gpt2",
+    max_new_tokens=200,
+    do_sample=False
+)
+
+def critique(text: str) -> str:
+    prompt = (
+        "Review this draft. Identify unclear reasoning, gaps, contradictions.\n\n"
+        "DRAFT:\n" + text + "\n\nReturn critique only:\n"
+    )
+    out = critic(prompt)[0]["generated_text"]
+    return out[len(prompt):].strip() if out.startswith(prompt) else out.strip()
+
+def refine(text: str, critique_text: str) -> str:
+    prompt = (
+        "Improve the draft using the critique. Fix gaps, strengthen logic.\n\n"
+        "CRITIQUE:\n" + critique_text +
+        "\n\nDRAFT:\n" + text +
+        "\n\nReturn improved output:\n"
+    )
+    out = critic(prompt)[0]["generated_text"]
+    return out[len(prompt):].strip() if out.startswith(prompt) else out.strip()
+
+def critique_and_refine(text: str) -> str:
+    c = critique(text)
+    return refine(text, c)
+
+
+# =============================================================
+# LLM CALL + PROMPT BUILDING
+# =============================================================
 def call_llm(prompt: str) -> str:
     out = llm(prompt, do_sample=False)[0]["generated_text"]
-    if out.startswith(prompt):
-        out = out[len(prompt):]
-    return out.strip()
+    return out[len(prompt):].strip() if out.startswith(prompt) else out.strip()
 
 
-# -------------------------------------------------------------
-# Build prompts
-# -------------------------------------------------------------
 def build_generator_prompt(module_id: str, *inputs: str) -> str:
     m = GEN_BY_ID[module_id]
     keys = list(m["input_placeholders"].keys())
@@ -61,21 +170,17 @@ def build_generator_prompt(module_id: str, *inputs: str) -> str:
     secs = m["output_sections"]
 
     p = []
-    p.append("You are a structured reasoning module.")
     p.append(f"MODULE: {m['label']} (id={module_id})")
-    p.append("")
+    p.append("You must follow the structured reasoning format.\n")
     p.append("INPUTS:")
     for k, v in vals.items():
         p.append(f"{k.upper()}: {v}")
-    p.append("")
-    p.append("You must respond using these sections:")
+    p.append("\nOutput sections:")
     for s in secs:
         p.append(f"- {s}")
-    p.append("")
+    p.append("\nFormat exactly as:")
     for s in secs:
-        p.append(f"{s}:")
-        p.append("[content]")
-        p.append("")
+        p.append(f"{s}:\n[content]\n")
     return "\n".join(p)
 
 
@@ -84,130 +189,100 @@ def build_checker_prompt(checker_id: str, *vals: str) -> str:
     secs = c["output_sections"]
 
     if len(vals) < 2:
-        original_task = ""
+        original = ""
         draft = vals[0] if vals else ""
     else:
-        original_task = "\n\n".join(vals[:-1])
+        original = "\n\n".join(vals[:-1])
         draft = vals[-1]
 
     p = []
-    p.append("You are a strict reviewer.")
     p.append(f"CHECKER: {c['label']} (id={checker_id})")
-    p.append("")
-    p.append("ORIGINAL TASK:")
-    p.append(original_task)
-    p.append("")
-    p.append("DRAFT OUTPUT:")
-    p.append(draft)
-    p.append("")
-    p.append("You must respond using:")
+    p.append("Review for structure, alignment and reasoning quality.\n")
+    p.append("ORIGINAL TASK:\n" + original + "\n")
+    p.append("DRAFT OUTPUT:\n" + draft + "\n")
+    p.append("Sections required:")
     for s in secs:
         p.append(f"- {s}")
-    p.append("")
+    p.append("\nFormat:")
     for s in secs:
-        p.append(f"{s}:")
-        p.append("[content]")
-        p.append("")
+        p.append(f"{s}:\n[content]\n")
     return "\n".join(p)
 
 
-# -------------------------------------------------------------
-# Generator & Checker Execution
-# -------------------------------------------------------------
+# =============================================================
+# GENERATOR + CHECKER EXECUTION
+# =============================================================
 def run_generator(module_id: str, *inputs: str) -> str:
     m = GEN_BY_ID[module_id]
 
-    # Load domain adapter if defined
     if m.get("domain"):
-        load_adapter(model, m["domain"])
+        load_domain_adapter(m["domain"])
 
-    # Build prompt
     prompt = build_generator_prompt(module_id, *inputs)
-
-    # Apply reasoning scaffolds (CoT + critique loop)
     prompt = apply_cot(prompt)
     draft = call_llm(prompt)
     final = critique_and_refine(draft)
-
     return final
 
 
 def run_checker(checker_id: str, *inputs: str) -> str:
     prompt = build_checker_prompt(checker_id, *inputs)
     prompt = apply_cot(prompt)
-    reviewed = call_llm(prompt)
-    return reviewed
-
-
-# -------------------------------------------------------------
-# Hybrid Router (rules + zero-shot)
-# -------------------------------------------------------------
-def hybrid_route(task_text: str):
-    if not task_text or not task_text.strip():
-        return "No task provided", "", ""
+    return call_llm(prompt)
 
-    # 1. Rule-based (deterministic)
-    rule_choice = rule_router(task_text)
-    if rule_choice:
-        return GEN_BY_ID[rule_choice]["label"], rule_choice, "Rule-based match"
 
-    # 2. Zero-shot fallback
-    predicted_label, module_id, scores = classify_task(task_text)
-    return predicted_label, module_id, scores
-
-
-# -------------------------------------------------------------
-# UI
-# -------------------------------------------------------------
+# =============================================================
+# GRADIO UI
+# =============================================================
 def build_ui():
-    with gr.Blocks(title="Modular Intelligence") as demo:
-        gr.Markdown("# Modular Intelligence\nUpgraded architecture with routing, adapters, and reasoning layers.")
+    with gr.Blocks(title="Modular Intelligence — Unified System") as demo:
+
+        gr.Markdown("# Modular Intelligence\nUnified architecture with routing, adapters, and reasoning scaffolds.")
 
-        # -------------------- Auto-Route Tab --------------------
+        # ---------------- AUTO-ROUTE TAB ----------------
         with gr.Tab("Auto-Route"):
             task_box = gr.Textbox(label="Describe your task", lines=6)
-            module_name = gr.Textbox(label="Suggested Module", interactive=False)
-            module_id = gr.Textbox(label="Module ID", interactive=False)
-            scores = gr.Textbox(label="Routing Details", lines=12, interactive=False)
+            out_name = gr.Textbox(label="Suggested Module", interactive=False)
+            out_id = gr.Textbox(label="Module ID", interactive=False)
+            out_scores = gr.Textbox(label="Routing Details", lines=12, interactive=False)
 
-            classify_btn = gr.Button("Classify Task")
-            classify_btn.click(
+            gr.Button("Classify Task").click(
                 fn=hybrid_route,
                 inputs=[task_box],
-                outputs=[module_name, module_id, scores],
+                outputs=[out_name, out_id, out_scores],
             )
 
-        # -------------------- Module Tabs ------------------------
+        # ---------------- MODULE TABS ----------------
         for m in GENERATORS:
             with gr.Tab(m["label"]):
-                gr.Markdown(f"**Module ID:** `{m['id']}` | **Domain:** `{m.get('domain', 'general')}`")
+                gr.Markdown(f"**Module ID:** `{m['id']}` | **Domain:** `{m.get('domain','general')}`")
 
                 inputs = []
                 for key, placeholder in m["input_placeholders"].items():
                     t = gr.Textbox(label=key, placeholder=placeholder, lines=4)
                     inputs.append(t)
 
-                output_box = gr.Textbox(label="Module Output", lines=20)
+                output = gr.Textbox(label="Generator Output", lines=18)
                 gr.Button("Run Module").click(
                     fn=partial(run_generator, m["id"]),
                     inputs=inputs,
-                    outputs=output_box,
+                    outputs=output,
                 )
 
                 checker_id = m.get("checker_id")
-                if checker_id and checker_id in CHECKERS:
-                    checker_output = gr.Textbox(label="Checker Output", lines=15)
+                if checker_id in CHECKERS:
+                    check_out = gr.Textbox(label="Checker Output", lines=15)
                     gr.Button("Run Checker").click(
                         fn=partial(run_checker, checker_id),
-                        inputs=inputs + [output_box],
-                        outputs=checker_output,
+                        inputs=inputs + [output],
+                        outputs=check_out,
                     )
                 else:
-                    gr.Markdown("_No checker available for this module._")
+                    gr.Markdown("_No checker for this module._")
 
     return demo
 
 
 if __name__ == "__main__":
-    app = build_ui()
-    app.launch()
+    ui = build_ui()
+    ui.launch()