demo_app.py

"""
Web Demo — Bengali Public Service Dialog System
=================================================
Simple Gradio-based chat interface for demonstrating the dialog system.
Can be run on Colab or locally.

Usage (Colab):
    !pip install -q gradio
    !python demo_app.py \
        --nlu_model_dir /content/models/joint_intent_ner \
        --gen_model_dir /content/models/response_gen/best_model \
        --labels_dir /content/data/processed/labels

Usage (local):
    python demo_app.py --nlu_model_dir ... --gen_model_dir ... --labels_dir ...
"""

import argparse
import json
import os
import sys
import uuid

import torch

# Ensure project modules are importable
script_dir = os.path.dirname(os.path.abspath(__file__))
project_dir = os.path.dirname(script_dir)
sys.path.insert(0, project_dir)
sys.path.insert(0, os.path.join(project_dir, "model"))
sys.path.insert(0, script_dir)

from pipeline import DialogPipeline


def create_demo(pipeline: DialogPipeline):
    """Build the Gradio chat interface."""
    import gradio as gr

    # Track conversations per session
    conversations = {}

    def chat(message, history, domain):
        # Get or create conversation ID
        conv_id = "demo_" + str(id(history)) if history else "demo_" + uuid.uuid4().hex[:8]

        if not history:
            pipeline.start_conversation(conv_id, domain)

        result = pipeline.respond(conv_id, message, domain_hint=domain)

        # Build info string
        info = (
            f"Intent: {result['intent']} ({result['confidence']:.2f}) | "
            f"Domain: {result['domain']} | "
            f"State: {result['state']}"
        )
        if result["entities"]:
            info += f" | Entities: {result['entities']}"
        if result["should_escalate"]:
            info += " | ⚠️ ESCALATION SUGGESTED"

        response = result["response"]
        if result["filled_slots"]:
            response += f"\n\n---\n🔍 {info}"
        else:
            response += f"\n\n---\n🔍 {info}"

        return response

    demo = gr.ChatInterface(
        fn=chat,
        additional_inputs=[
            gr.Dropdown(
                choices=["general", "passport", "nid", "utilities", "welfare"],
                value="general",
                label="Service Domain / সেবার ধরন",
            ),
        ],
        title="🇧🇩 Bengali Public Service Dialog System",
        description=(
            "বাংলাদেশ সরকারি সেবা সহায়তা ব্যবস্থা\n\n"
            "Ask questions about passport, NID, utilities, welfare services "
            "in Bengali (Standard, Sylheti, or Chittagonian dialect)."
        ),
        examples=[
            ["আমি পাসপোর্ট করতে চাই।", "passport"],
            ["বিদ্যুৎ বিল কীভাবে দেব?", "utilities"],
            ["বয়স্ক ভাতার জন্য আবেদন করতে চাই।", "welfare"],
            ["হামি এনআইডি কার্ড বানাইতে চাই।", "nid"],
            ["আঁই পাসপোর্ট বানাইত্তে সাই।", "passport"],
        ],
    )

    return demo


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--nlu_model_dir", type=str, required=True)
    parser.add_argument("--gen_model_dir", type=str, required=True)
    parser.add_argument("--labels_dir", type=str, required=True)
    parser.add_argument("--bert_model", type=str,
                        default="sagorsarker/bangla-bert-base")
    parser.add_argument("--port", type=int, default=7860)
    parser.add_argument("--share", action="store_true",
                        help="Create public Gradio link")
    args = parser.parse_args()

    # Load pipeline
    pipeline = DialogPipeline(
        nlu_model_path=args.nlu_model_dir,
        gen_model_path=args.gen_model_dir,
        labels_dir=args.labels_dir,
        bert_model=args.bert_model,
    )

    # Create and launch demo
    demo = create_demo(pipeline)
    demo.launch(server_port=args.port, share=args.share)


if __name__ == "__main__":
    main()