init

Signed-off-by: bitliu <[email protected]>

app.py

@@ -1,7 +1,11 @@
 import streamlit as st
 import streamlit.components.v1 as components
 import torch
-from transformers import AutoTokenizer, AutoModelForSequenceClassification, AutoModelForTokenClassification
+from transformers import (
+    AutoTokenizer,
+    AutoModelForSequenceClassification,
+    AutoModelForTokenClassification,
+)
 
 # ============== Model Configurations ==============
 MODELS = {
@@ -10,11 +14,20 @@ MODELS = {
         "description": "Classifies prompts into academic/professional categories.",
         "type": "sequence",
         "labels": {
-            0: ("biology", "🧬"), 1: ("business", "💼"), 2: ("chemistry", "🧪"),
-            3: ("computer science", "💻"), 4: ("economics", "📈"), 5: ("engineering", "⚙️"),
-            6: ("health", "🏥"), 7: ("history", "📜"), 8: ("law", "⚖️"),
-            9: ("math", "🔢"), 10: ("other", "📦"), 11: ("philosophy", "🤔"),
-            12: ("physics", "⚛️"), 13: ("psychology", "🧠"),
+            0: ("biology", "🧬"),
+            1: ("business", "💼"),
+            2: ("chemistry", "🧪"),
+            3: ("computer science", "💻"),
+            4: ("economics", "📈"),
+            5: ("engineering", "⚙️"),
+            6: ("health", "🏥"),
+            7: ("history", "📜"),
+            8: ("law", "⚖️"),
+            9: ("math", "🔢"),
+            10: ("other", "📦"),
+            11: ("philosophy", "🤔"),
+            12: ("physics", "⚛️"),
+            13: ("psychology", "🧠"),
         },
         "demo": "What is photosynthesis and how does it work?",
     },
@@ -37,12 +50,24 @@ MODELS = {
         "description": "Detects the primary type of PII in the text.",
         "type": "sequence",
         "labels": {
-            0: ("AGE", "🎂"), 1: ("CREDIT_CARD", "💳"), 2: ("DATE_TIME", "📅"),
-            3: ("DOMAIN_NAME", "🌐"), 4: ("EMAIL_ADDRESS", "📧"), 5: ("GPE", "🗺️"),
-            6: ("IBAN_CODE", "🏦"), 7: ("IP_ADDRESS", "🖥️"), 8: ("NO_PII", "✅"),
-            9: ("NRP", "👥"), 10: ("ORGANIZATION", "🏢"), 11: ("PERSON", "👤"),
-            12: ("PHONE_NUMBER", "📞"), 13: ("STREET_ADDRESS", "🏠"), 14: ("TITLE", "📛"),
-            15: ("US_DRIVER_LICENSE", "🚗"), 16: ("US_SSN", "🔐"), 17: ("ZIP_CODE", "📮"),
+            0: ("AGE", "🎂"),
+            1: ("CREDIT_CARD", "💳"),
+            2: ("DATE_TIME", "📅"),
+            3: ("DOMAIN_NAME", "🌐"),
+            4: ("EMAIL_ADDRESS", "📧"),
+            5: ("GPE", "🗺️"),
+            6: ("IBAN_CODE", "🏦"),
+            7: ("IP_ADDRESS", "🖥️"),
+            8: ("NO_PII", "✅"),
+            9: ("NRP", "👥"),
+            10: ("ORGANIZATION", "🏢"),
+            11: ("PERSON", "👤"),
+            12: ("PHONE_NUMBER", "📞"),
+            13: ("STREET_ADDRESS", "🏠"),
+            14: ("TITLE", "📛"),
+            15: ("US_DRIVER_LICENSE", "🚗"),
+            16: ("US_SSN", "🔐"),
+            17: ("ZIP_CODE", "📮"),
         },
         "demo": "My email is [email protected] and my phone is 555-123-4567",
     },
@@ -53,6 +78,32 @@ MODELS = {
         "labels": None,
         "demo": "John Smith works at Microsoft in Seattle, his email is [email protected]",
     },
+    "😤 Dissatisfaction Detector": {
+        "id": "llm-semantic-router/dissat-detector",
+        "description": "Detects user dissatisfaction in conversational AI interactions. Classifies user follow-up messages as satisfied (SAT) or dissatisfied (DISSAT).",
+        "type": "dialogue",
+        "labels": {0: ("SAT", "🟢"), 1: ("DISSAT", "🔴")},
+        "demo": {
+            "query": "Find a restaurant nearby",
+            "response": "I found Italian Kitchen for you.",
+            "followup": "Show me other options",
+        },
+    },
+    "🔍 Dissatisfaction Explainer": {
+        "id": "llm-semantic-router/dissat-explainer",
+        "description": "Explains why a user is dissatisfied. Stage 2 of hierarchical dissatisfaction detection - classifies into NEED_CLARIFICATION, WRONG_ANSWER, or WANT_DIFFERENT.",
+        "type": "dialogue",
+        "labels": {
+            0: ("NEED_CLARIFICATION", "❓"),
+            1: ("WRONG_ANSWER", "❌"),
+            2: ("WANT_DIFFERENT", "🔄"),
+        },
+        "demo": {
+            "query": "Book a table for 2",
+            "response": "Table for 3 confirmed",
+            "followup": "No, I said 2 people not 3",
+        },
+    },
 }
 
 
@@ -78,7 +129,29 @@ def classify_sequence(text: str, model_id: str, labels: dict) -> tuple:
     pred_class = torch.argmax(probs).item()
     label_name, emoji = labels[pred_class]
     confidence = probs[pred_class].item()
-    all_scores = {f"{labels[i][1]} {labels[i][0]}": float(probs[i]) for i in range(len(labels))}
+    all_scores = {
+        f"{labels[i][1]} {labels[i][0]}": float(probs[i]) for i in range(len(labels))
+    }
+    return label_name, emoji, confidence, all_scores
+
+
+def classify_dialogue(
+    query: str, response: str, followup: str, model_id: str, labels: dict
+) -> tuple:
+    """Classify dialogue using sequence classification model with special format."""
+    tokenizer, model = load_model(model_id, "sequence")
+    # Format input as per model requirements
+    text = f"[USER QUERY] {query}\n[SYSTEM RESPONSE] {response}\n[USER FOLLOWUP] {followup}"
+    inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=512)
+    with torch.no_grad():
+        outputs = model(**inputs)
+        probs = torch.softmax(outputs.logits, dim=-1)[0]
+    pred_class = torch.argmax(probs).item()
+    label_name, emoji = labels[pred_class]
+    confidence = probs[pred_class].item()
+    all_scores = {
+        f"{labels[i][1]} {labels[i][0]}": float(probs[i]) for i in range(len(labels))
+    }
     return label_name, emoji, confidence, all_scores
 
 
@@ -86,7 +159,13 @@ def classify_tokens(text: str, model_id: str) -> list:
     """Token-level NER classification."""
     tokenizer, model = load_model(model_id, "token")
     id2label = model.config.id2label
-    inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=512, return_offsets_mapping=True)
+    inputs = tokenizer(
+        text,
+        return_tensors="pt",
+        truncation=True,
+        max_length=512,
+        return_offsets_mapping=True,
+    )
     offset_mapping = inputs.pop("offset_mapping")[0].tolist()
     with torch.no_grad():
         outputs = model(**inputs)
@@ -101,7 +180,11 @@ def classify_tokens(text: str, model_id: str) -> list:
             if current_entity:
                 entities.append(current_entity)
             current_entity = {"type": label[2:], "start": start, "end": end}
-        elif label.startswith("I-") and current_entity and label[2:] == current_entity["type"]:
+        elif (
+            label.startswith("I-")
+            and current_entity
+            and label[2:] == current_entity["type"]
+        ):
             current_entity["end"] = end
         else:
             if current_entity:
@@ -110,7 +193,7 @@ def classify_tokens(text: str, model_id: str) -> list:
     if current_entity:
         entities.append(current_entity)
     for e in entities:
-        e["text"] = text[e["start"]:e["end"]]
+        e["text"] = text[e["start"] : e["end"]]
     return entities
 
 
@@ -119,13 +202,21 @@ def create_highlighted_html(text: str, entities: list) -> str:
     if not entities:
         return f'<div style="padding:15px;background:#f0f0f0;border-radius:8px;">{text}</div>'
     html = text
-    colors = {"EMAIL_ADDRESS": "#ff6b6b", "PHONE_NUMBER": "#4ecdc4", "PERSON": "#45b7d1",
-              "STREET_ADDRESS": "#96ceb4", "US_SSN": "#d63384", "CREDIT_CARD": "#fd7e14",
-              "ORGANIZATION": "#6f42c1", "GPE": "#20c997", "IP_ADDRESS": "#0dcaf0"}
+    colors = {
+        "EMAIL_ADDRESS": "#ff6b6b",
+        "PHONE_NUMBER": "#4ecdc4",
+        "PERSON": "#45b7d1",
+        "STREET_ADDRESS": "#96ceb4",
+        "US_SSN": "#d63384",
+        "CREDIT_CARD": "#fd7e14",
+        "ORGANIZATION": "#6f42c1",
+        "GPE": "#20c997",
+        "IP_ADDRESS": "#0dcaf0",
+    }
     for e in sorted(entities, key=lambda x: x["start"], reverse=True):
         color = colors.get(e["type"], "#ffc107")
         span = f'<span style="background:{color};padding:2px 6px;border-radius:4px;color:white;" title="{e["type"]}">{e["text"]}</span>'
-        html = html[:e["start"]] + span + html[e["end"]:]
+        html = html[: e["start"]] + span + html[e["end"] :]
     return f'<div style="padding:15px;background:#f8f9fa;border-radius:8px;line-height:2;">{html}</div>'
 
 
@@ -135,10 +226,15 @@ def main():
     # Header with logo
     col1, col2 = st.columns([1, 4])
     with col1:
-        st.image("https://github.com/vllm-project/semantic-router/blob/main/website/static/img/vllm.png?raw=true", width=150)
+        st.image(
+            "https://github.com/vllm-project/semantic-router/blob/main/website/static/img/vllm.png?raw=true",
+            width=150,
+        )
     with col2:
         st.title("🧠 LLM Semantic Router")
-        st.markdown("**Intelligent Router for Mixture-of-Models** | Part of the [vLLM](https://github.com/vllm-project/vllm) ecosystem")
+        st.markdown(
+            "**Intelligent Router for Mixture-of-Models** | Part of the [vLLM](https://github.com/vllm-project/vllm) ecosystem"
+        )
 
     st.markdown("---")
 
@@ -161,18 +257,70 @@ def main():
 
     # Main content
     st.subheader("📝 Input")
-    text_input = st.text_area(
-        "Enter text to analyze:",
-        value=model_config["demo"],
-        height=120,
-        placeholder="Type your text here..."
-    )
+
+    # Different input UI based on model type
+    if model_config["type"] == "dialogue":
+        # Dialogue models need query, response, and followup
+        demo = model_config["demo"]
+        query_input = st.text_input(
+            "🗣️ User Query:",
+            value=demo["query"],
+            placeholder="Enter the original user query...",
+        )
+        response_input = st.text_input(
+            "🤖 System Response:",
+            value=demo["response"],
+            placeholder="Enter the system's response...",
+        )
+        followup_input = st.text_input(
+            "💬 User Follow-up:",
+            value=demo["followup"],
+            placeholder="Enter the user's follow-up message...",
+        )
+        text_input = None  # Not used for dialogue models
+    else:
+        # Standard text input for other models
+        text_input = st.text_area(
+            "Enter text to analyze:",
+            value=model_config["demo"],
+            height=120,
+            placeholder="Type your text here...",
+        )
+        query_input = response_input = followup_input = None
 
     st.markdown("---")
 
     # Analyze button
     if st.button("🔍 Analyze", type="primary", use_container_width=True):
-        if not text_input.strip():
+        if model_config["type"] == "dialogue":
+            if (
+                not query_input.strip()
+                or not response_input.strip()
+                or not followup_input.strip()
+            ):
+                st.warning("Please fill in all dialogue fields.")
+            else:
+                with st.spinner("Analyzing..."):
+                    label, emoji, conf, scores = classify_dialogue(
+                        query_input,
+                        response_input,
+                        followup_input,
+                        model_config["id"],
+                        model_config["labels"],
+                    )
+                    st.session_state.result = {
+                        "type": "dialogue",
+                        "label": label,
+                        "emoji": emoji,
+                        "confidence": conf,
+                        "scores": scores,
+                        "input": {
+                            "query": query_input,
+                            "response": response_input,
+                            "followup": followup_input,
+                        },
+                    }
+        elif not text_input.strip():
             st.warning("Please enter some text to analyze.")
         else:
             with st.spinner("Analyzing..."):
@@ -185,14 +333,14 @@ def main():
                         "label": label,
                         "emoji": emoji,
                         "confidence": conf,
-                        "scores": scores
+                        "scores": scores,
                     }
                 else:
                     entities = classify_tokens(text_input, model_config["id"])
                     st.session_state.result = {
                         "type": "token",
                         "entities": entities,
-                        "text": text_input
+                        "text": text_input,
                     }
 
     # Display results
@@ -200,24 +348,28 @@ def main():
         st.markdown("---")
         st.subheader("📊 Results")
         result = st.session_state.result
-        if result["type"] == "sequence":
+        if result["type"] in ("sequence", "dialogue"):
             col1, col2 = st.columns([1, 1])
             with col1:
                 st.success(f"{result['emoji']} **{result['label']}**")
                 st.metric("Confidence", f"{result['confidence']:.1%}")
             with col2:
                 st.markdown("**All Scores:**")
-                sorted_scores = dict(sorted(result["scores"].items(), key=lambda x: x[1], reverse=True))
+                sorted_scores = dict(
+                    sorted(result["scores"].items(), key=lambda x: x[1], reverse=True)
+                )
                 for k, v in sorted_scores.items():
                     st.progress(v, text=f"{k}: {v:.1%}")
-        else:
+        elif result["type"] == "token":
             entities = result["entities"]
             if entities:
                 st.success(f"Found {len(entities)} PII entity(s)")
                 for e in entities:
                     st.markdown(f"- **{e['type']}**: `{e['text']}`")
                 st.markdown("### Highlighted Text")
-                components.html(create_highlighted_html(result["text"], entities), height=150)
+                components.html(
+                    create_highlighted_html(result["text"], entities), height=150
+                )
             else:
                 st.info("✅ No PII detected")
 
@@ -235,9 +387,9 @@ def main():
         <b>GitHub</b>: <a href="https://github.com/vllm-project/semantic-router">vllm-project/semantic-router</a>
         </div>
         """,
-        unsafe_allow_html=True
+        unsafe_allow_html=True,
     )
 
 
 if __name__ == "__main__":
-    main()
+    main()