Initial commit of Arch Router Simulator

README.md

@@ -1,8 +1,8 @@
 ---
 title: ArchRouter Simulator
-emoji: 🔥
-colorFrom: gray
-colorTo: indigo
+emoji: 🧭
+colorFrom: yellow
+colorTo: blue
 sdk: gradio
 sdk_version: 5.35.0
 app_file: app.py
@@ -10,4 +10,4 @@ pinned: false
 short_description: Demo working simulation of Arch Router
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+This is a demo for Arch Router

app.py

@@ -0,0 +1,208 @@
+import json
+import time
+import ast
+import torch
+import gradio as gr
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from gradio_consilium_roundtable import consilium_roundtable
+
+# === Constants ===
+MODEL_NAME = "katanemo/Arch-Router-1.5B"
+ARCH_ROUTER = "Arch Router"
+WAIT_DEPARTMENT = 5
+WAIT_SYSTEM = 5
+
+# === Load model/tokenizer ===
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_NAME, device_map="auto", torch_dtype="auto", trust_remote_code=True
+)
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+
+# === Route Definitions ===
+route_config = [
+    {"name": "code_generation", "description": "Generating code based on prompts"},
+    {"name": "bug_fixing", "description": "Fixing errors or bugs in code"},
+    {"name": "performance_optimization", "description": "Improving code performance"},
+    {"name": "api_help", "description": "Assisting with APIs and libraries"},
+    {"name": "programming", "description": "General programming Q&A"},
+    {"name": "legal", "description": "Legal"},
+    {"name": "healthcare", "description": "Healthcare and medical related"},
+]
+
+departments = {
+    "code_generation": ("💻", "Code Generation"),
+    "bug_fixing": ("🐞", "Bug Fixing"),
+    "performance_optimization": ("⚡", "Performance Optimization"),
+    "api_help": ("🔌", "API Help"),
+    "programming": ("📚", "Programming"),
+    "legal": ("⚖️", "Legal"),
+    "healthcare": ("🩺", "Healthcare"),
+    "other": ("❓", "Other / General Inquiry"),
+}
+
+# === Prompt Formatting ===
+TASK_INSTRUCTION = """
+You are a helpful assistant designed to find the best suited route. You are provided with route description within <routes></routes> XML tags:
+<routes>
+{routes}
+</routes>
+
+<conversation>
+{conversation}
+</conversation>
+"""
+
+FORMAT_PROMPT = """
+Your task is to decide which route is best suit with user intent on the conversation in <conversation></conversation> XML tags. Follow the instruction:
+1. If the latest intent from user is irrelevant or user intent is full filled, response with other route {"route": "other"}.
+2. You must analyze the route descriptions and find the best match route for user latest intent.
+3. You only response the name of the route that best matches the user's request, use the exact name in the <routes></routes>.
+
+Based on your analysis, provide your response in the following JSON format:
+{"route": "route_name"} 
+"""
+
+def format_prompt(route_config, conversation):
+    return TASK_INSTRUCTION.format(
+        routes=json.dumps(route_config), conversation=json.dumps(conversation)
+    ) + FORMAT_PROMPT
+
+def parse_route(response_text):
+    try:
+        start = response_text.find("{")
+        end = response_text.rfind("}") + 1
+        return ast.literal_eval(response_text[start:end]).get("route", "other")
+    except Exception as e:
+        print("Parsing failed:", e)
+        return "other"
+
+def init_state():
+    avatar_emojis = {
+        ARCH_ROUTER: "https://avatars.githubusercontent.com/u/112724757?s=200&v=4",
+        "code_generation": "https://raw.githubusercontent.com/twitter/twemoji/master/assets/72x72/1f4bb.png",
+        "bug_fixing": "https://raw.githubusercontent.com/twitter/twemoji/master/assets/72x72/1f41e.png",
+        "performance_optimization": "https://raw.githubusercontent.com/twitter/twemoji/master/assets/72x72/26a1.png",
+        "api_help": "https://raw.githubusercontent.com/twitter/twemoji/master/assets/72x72/1f50c.png",
+        "programming": "https://raw.githubusercontent.com/twitter/twemoji/master/assets/72x72/1f4da.png",
+        "legal": "https://raw.githubusercontent.com/twitter/twemoji/master/assets/72x72/2696.png",
+        "healthcare": "https://raw.githubusercontent.com/twitter/twemoji/master/assets/72x72/1fa7a.png",
+        "other": "https://raw.githubusercontent.com/twitter/twemoji/master/assets/72x72/2753.png",
+    }
+    return {
+        "messages": [],
+        "participants": [ARCH_ROUTER] + list(departments.keys()),
+        "currentSpeaker": None,
+        "thinking": [],
+        "showBubbles": [ARCH_ROUTER],
+        "avatarImages": avatar_emojis,
+    }
+
+def route_and_visualize(user_input_text, rt_state, chat_history):
+    chat_history = chat_history or []
+    rt_state = rt_state or {"messages": []}
+    chat_history.append(("User", user_input_text))
+
+    # Step 1: Disable input and show route detection
+    rt_state["messages"] = [{"speaker": ARCH_ROUTER, "text": "🔎 Identifying route, please wait..."}]
+    yield rt_state, chat_history, rt_state, gr.update(interactive=False)
+
+    # Step 2: Prepare prompt and get route
+    conversation = [{"role": "user", "content": user_input_text}]
+    route_prompt = format_prompt(route_config, conversation)
+    input_ids = tokenizer.apply_chat_template(
+        [{"role": "user", "content": route_prompt}],
+        add_generation_prompt=True,
+        return_tensors="pt"
+    ).to(model.device)
+
+    with torch.no_grad():
+        output = model.generate(input_ids=input_ids, max_new_tokens=512)
+
+    prompt_len = input_ids.shape[1]
+    response = tokenizer.decode(output[0][prompt_len:], skip_special_tokens=True).strip()
+    print("MODEL RAW:", response)
+    route = parse_route(response)
+
+    emoji, dept_name = departments.get(route, departments["other"])
+
+    # Step 3: Show route identified
+    rt_state["messages"][0] = {
+        "speaker": ARCH_ROUTER,
+        "text": f"📌 Identified department: **{dept_name}**. Forwarding task..."
+    }
+    chat_history.append((ARCH_ROUTER, f"📌 Identified department: {dept_name}. Forwarding task..."))
+    yield rt_state, chat_history, rt_state, gr.update(interactive=False)
+
+    # Step 4: Show processing
+    time.sleep(3)
+    rt_state["messages"].extend([
+        {"speaker": route, "text": f"{emoji} {dept_name} simulation is processing your request in {WAIT_DEPARTMENT} secs..."},
+        {"speaker": ARCH_ROUTER, "text": "⏳ Waiting for department to respond..."}
+    ])
+    rt_state["showBubbles"] = [ARCH_ROUTER, route]
+    yield rt_state, chat_history, rt_state, gr.update(interactive=False)
+
+    # Step 5: Simulate delay and complete
+    time.sleep(WAIT_DEPARTMENT)
+    rt_state["messages"][-2]["text"] = f"✅ {dept_name} completed the task."
+    rt_state["messages"][-1]["text"] = f"✅ {dept_name} department has completed the task."
+    chat_history.append((ARCH_ROUTER, f"✅ {dept_name} department completed the task."))
+    yield rt_state, chat_history, rt_state, gr.update(interactive=False)
+
+    # Step 6: Reset visible bubbles
+    rt_state["showBubbles"] = [ARCH_ROUTER]
+    yield rt_state, chat_history, rt_state, gr.update(interactive=False)
+
+    # Step 7: System ready
+    time.sleep(WAIT_SYSTEM)
+    rt_state["messages"].append({"speaker": ARCH_ROUTER, "text": "Arch Router is ready to discuss."})
+    yield rt_state, chat_history, rt_state, gr.update(interactive=True)
+
+# === Gradio UI ===
+with gr.Blocks(theme=gr.themes.Ocean()) as demo:
+    gr.Markdown("## 🧭 Arch Router Simulation: Smart Department Dispatcher")
+
+    gr.Markdown(
+        """
+        <p>This is a demo simulation of <a href="https://huggingface.co/katanemo/Arch-Router-1.5B" target="_blank">katanemo/Arch-Router-1.5B</a></p>
+        """
+    )
+
+    with gr.Row():
+        with gr.Column(scale=3):
+            rt_state = gr.State(init_state())
+            chat_state = gr.State([])
+            roundtable = consilium_roundtable(value=init_state())
+
+        with gr.Column(scale=2):
+            chatbot = gr.Chatbot(label="Chat History", max_height=300)
+            textbox = gr.Textbox(placeholder="Describe your issue...", label="Ask Arch Router")
+            submit_btn = gr.Button("Submit")
+
+            example_inputs = [
+                "How do I optimize this loop in Python?",
+                "Generate a function to sort an array in python",
+                "Help me anonymize patient health records before storing them",
+                "I'm getting a TypeError in following code",
+                "Do I need to include attribution for MIT-licensed software?",
+                "How do I connect to external API from this code?"
+            ]
+
+            # Trigger submission via Enter or Button
+            for trigger in (textbox.submit, submit_btn.click):
+                trigger(
+                    route_and_visualize,
+                    inputs=[textbox, rt_state, chat_state],
+                    outputs=[roundtable, chatbot, rt_state, textbox],
+                    concurrency_limit=1
+                )
+
+            # Example block
+            gr.Examples(
+                examples=example_inputs,
+                inputs=textbox,
+                label="Try one of these examples"
+            )
+
+if __name__ == "__main__":
+    demo.launch()

requirements.txt

@@ -0,0 +1,6 @@
+transformers>=4.30.0
+torch>=1.13.0
+gradio>=3.0
+accelerate
+sentencepiece
+gradio_consilium_roundtable