Update app.py

app.py

@@ -1,70 +1,143 @@
+import re
+import numpy as np
+from typing import Dict
+
 import gradio as gr
-from transformers import pipeline
+from fastapi import FastAPI
+from sentence_transformers import SentenceTransformer
+from sklearn.metrics.pairwise import cosine_similarity
 
-# Load zero-shot classifier
-classifier = pipeline(
-    "zero-shot-classification",
-    model="valhalla/distilbart-mnli-12-3"
-)
+# -------------------------
+# CONFIG
+# -------------------------
+
+MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2"
+CONFIDENCE_THRESHOLD = 0.35
+
+# -------------------------
+# LOAD MODEL
+# -------------------------
+
+embedder = SentenceTransformer(MODEL_NAME)
+
+# -------------------------
+# RULE-BASED ROUTER
+# -------------------------
+
+GREETINGS = {
+    "hi", "hello", "hey", "yo", "sup", "hola", "hii", "hai"
+}
+
+IMAGE_KEYWORDS = {
+    "draw", "image", "picture", "photo", "generate image", "create image"
+}
+
+CODE_KEYWORDS = {
+    "code", "python", "javascript", "bug", "error", "compile", "program"
+}
+
+def rule_based_route(text: str):
+    t = text.lower().strip()
+
+    if t in GREETINGS:
+        return "chat"
 
-# Default routing labels (edit freely)
-DEFAULT_LABELS = [
-    "chat",
-    "search",
-    "image_generation",
-    "code",
-    "research",
-    "study",
-    "project",
-    "action"
-]
+    if any(k in t for k in IMAGE_KEYWORDS):
+        return "image_generation"
 
-def classify(text, labels_text):
-    if not text.strip():
-        return {}
+    if any(k in t for k in CODE_KEYWORDS):
+        return "code"
 
-    labels = [l.strip() for l in labels_text.split(",") if l.strip()]
-    if not labels:
-        labels = DEFAULT_LABELS
+    return None  # fallback to semantic router
 
-    result = classifier(text, labels, multi_label=False)
+# -------------------------
+# SEMANTIC INTENTS
+# -------------------------
 
+INTENTS: Dict[str, str] = {
+    "chat": "casual conversation, greetings, talking",
+    "search": "asking for information or facts",
+    "image_generation": "requesting image creation or visual generation",
+    "code": "programming, software development, debugging",
+    "research": "deep technical or academic research",
+    "study": "learning, studying, explanations, tutorials",
+    "project": "building, planning, or creating a project",
+    "action": "asking the system to perform an action"
+}
+
+INTENT_NAMES = list(INTENTS.keys())
+INTENT_EMBEDDINGS = embedder.encode(list(INTENTS.values()), normalize_embeddings=True)
+
+# -------------------------
+# SEMANTIC ROUTER
+# -------------------------
+
+def semantic_route(text: str):
+    text_emb = embedder.encode([text], normalize_embeddings=True)
+    sims = cosine_similarity(text_emb, INTENT_EMBEDDINGS)[0]
+
+    scores = dict(zip(INTENT_NAMES, sims))
+    top_intent = max(scores, key=scores.get)
+    confidence = scores[top_intent]
+
+    if confidence < CONFIDENCE_THRESHOLD:
+        return "chat", scores  # safe fallback
+
+    return top_intent, scores
+
+# -------------------------
+# MAIN ROUTER
+# -------------------------
+
+def route_intent(text: str):
+    # 1️⃣ Rule-based (instant)
+    rule = rule_based_route(text)
+    if rule:
+        return {
+            "text": text,
+            "top_intent": rule,
+            "method": "rule",
+            "scores": {rule: 1.0}
+        }
+
+    # 2️⃣ Semantic
+    intent, scores = semantic_route(text)
     return {
         "text": text,
-        "top_intent": result["labels"][0],
-        "scores": dict(zip(result["labels"], result["scores"]))
+        "top_intent": intent,
+        "method": "semantic",
+        "scores": scores
     }
 
-# -------- Gradio UI --------
-with gr.Blocks(title="Zero-Shot Router (BART-MNLI)") as demo:
-    gr.Markdown("## 🧠 Zero-Shot Intent Router")
-    gr.Markdown(
-        "Classifies **any input** into your routing labels.\n"
-        "Used for **system-prompt injection + MPC routing**."
-    )
+# -------------------------
+# FASTAPI
+# -------------------------
 
-    text_input = gr.Textbox(
-        label="User Input",
-        placeholder="e.g. Generate a cyberpunk city wallpaper in 4k"
-    )
+api = FastAPI()
 
-    labels_input = gr.Textbox(
-        label="Labels (comma-separated)",
-        value="chat, search, image_generation, code, research, study, project, action"
-    )
+@api.post("/classify")
+def classify(payload: Dict):
+    text = payload.get("text", "")
+    return route_intent(text)
 
-    output = gr.JSON(label="Classification Result")
+# -------------------------
+# GRADIO UI
+# -------------------------
 
-    classify_btn = gr.Button("Classify")
+def gradio_classify(text):
+    result = route_intent(text)
+    return result
 
-    classify_btn.click(
-        fn=classify,
-        inputs=[text_input, labels_input],
-        outputs=output
-    )
+gradio_ui = gr.Interface(
+    fn=gradio_classify,
+    inputs=gr.Textbox(label="User Input"),
+    outputs=gr.JSON(label="Classification Result"),
+    title="🧠 Hybrid Intent Router",
+    description="Rule-based + Semantic intent classification for prompt routing & MPC selection"
+)
 
-# Expose BOTH UI + API
-app = demo
+# -------------------------
+# MOUNT GRADIO INTO FASTAPI
+# -------------------------
 
-if __name__ == "__main__":
-    demo.launch()
+app = gr.mount_gradio_app(api, gradio_ui, path="/")