src/chatbot.py

"""End-to-end multilingual chatbot pipeline (with generative RAG).

Architecture:
  user_text
    -> Language Detector  (AR / CS / EN / FR)            [CPU, DistilBERT]
    -> Intent Classifier  (booking / complaint / ...)    [CPU, DistilBERT]
    -> NER                (PER / LOC / ORG / DATE)       [CPU, DistilBERT]
    -> Branch on intent:
         greeting / farewell -> canned reply (deterministic, no LLM call)
         inquiry / booking / complaint -> RAG retrieve top-3 (same-lang preferred)
                                          + Qwen2.5-0.5B-Instruct generates reply
         other -> Qwen2.5-0.5B-Instruct answers from general knowledge (no retrieval)

Why this layout: the GTX 1650 has 4 GB VRAM. Putting the small DistilBERT
classifiers on CPU is essentially free (each <100 ms per turn) and frees the
GPU for the generator, which is the bottleneck.

Models loaded:
  models/lang_detector/      (DistilBERT, 4 classes)            [CPU]
  models/intent_classifier/  (DistilBERT, 6 classes)            [CPU]
  models/ner_model/          (DistilBERT TokenClassification)   [CPU]
  models/rag/faiss.index     (FAISS IndexFlatIP, 36 x 384)      [CPU]
  paraphrase-multilingual-MiniLM-L12-v2                          [CPU]
  Qwen/Qwen2.5-0.5B-Instruct (causal LM, fp16 on GPU)            [GPU]

Usage:
  python src/chatbot.py "How do I book a flight?"
  python src/chatbot.py "كيف أحجز فندقاً في باريس يوم 15 يونيو؟"
  python src/chatbot.py --interactive
  python src/chatbot.py --json "Hello there!"
"""

from __future__ import annotations

import argparse
import json
import sys
from dataclasses import dataclass, field, asdict
from pathlib import Path
from threading import Thread
from typing import Any, Iterator

import faiss
import numpy as np
import torch
from sentence_transformers import SentenceTransformer
from transformers import (
    AutoModelForCausalLM,
    AutoModelForSequenceClassification,
    AutoModelForTokenClassification,
    AutoTokenizer,
    TextIteratorStreamer,
    pipeline,
)

PROJECT_ROOT = Path(__file__).resolve().parent.parent


def _model_id(local_path: Path, hub_id: str) -> str:
    """Return local path if the model is trained locally, else the HF Hub repo id.
    Lets the same code run unchanged in dev (local weights) and in deployed
    environments like HF Spaces (weights pulled from the Hub).
    """
    return str(local_path) if local_path.exists() and any(local_path.iterdir()) else hub_id


LANG_DIR = _model_id(
    PROJECT_ROOT / "models" / "lang_detector",
    "momenalhamza/multilingual-chatbot-lang-detector",
)
INTENT_DIR = _model_id(
    PROJECT_ROOT / "models" / "intent_classifier",
    "momenalhamza/multilingual-chatbot-intent",
)
NER_DIR = _model_id(
    PROJECT_ROOT / "models" / "ner_model",
    "momenalhamza/multilingual-chatbot-ner",
)
RAG_DIR = PROJECT_ROOT / "models" / "rag"   # small (~80 KB), always shipped with the repo
EMBED_MODEL = "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"
GEN_MODEL = "Qwen/Qwen2.5-0.5B-Instruct"

LANG_NAME = {"AR": "Arabic", "EN": "English", "FR": "French",
             "CS": "Code-Switching (Arabic-English mix)"}

CANNED: dict[str, dict[str, str]] = {
    "greeting": {
        "AR": "مرحباً! كيف يمكنني مساعدتك اليوم؟",
        "EN": "Hello! How can I help you today?",
        "FR": "Bonjour ! Comment puis-je vous aider aujourd'hui ?",
        "CS": "Hello! How can I help you today?",
    },
    "farewell": {
        "AR": "إلى اللقاء! يومٌ سعيد.",
        "EN": "Goodbye! Have a great day.",
        "FR": "Au revoir ! Bonne journée.",
        "CS": "Goodbye! Have a great day.",
    },
    "fallback": {
        "AR": "عذراً، لم أفهم طلبك. هل يمكنك إعادة صياغته؟",
        "EN": "Sorry, I didn't understand. Could you rephrase?",
        "FR": "Désolé, je n'ai pas compris. Pourriez-vous reformuler ?",
        "CS": "Sorry, I didn't understand. Could you rephrase?",
    },
}


@dataclass
class Turn:
    """One chatbot exchange. Serializable to JSON for the UI/logs."""
    text: str
    language: str
    intent: str
    intent_confidence: float
    entities: list[dict]
    reply: str
    retrieved: list[dict] = field(default_factory=list)
    generator_used: bool = False


class Chatbot:
    """Holds all models. Thread-unsafe; create one per process."""

    def __init__(self, gen_max_new_tokens: int = 200) -> None:
        gpu_available = torch.cuda.is_available()
        gen_device = "cuda" if gpu_available else "cpu"
        print(f"[chatbot] GPU available: {gpu_available} -> generator on {gen_device}, "
              "classifiers on CPU")

        # Classifiers on CPU (free up VRAM for Qwen).
        print("[chatbot] loading language detector (CPU) ...")
        self.lang_pipe = pipeline(
            "text-classification",
            model=AutoModelForSequenceClassification.from_pretrained(str(LANG_DIR)),
            tokenizer=AutoTokenizer.from_pretrained(str(LANG_DIR)),
            device=-1, top_k=None,
        )

        print("[chatbot] loading intent classifier (CPU) ...")
        self.intent_pipe = pipeline(
            "text-classification",
            model=AutoModelForSequenceClassification.from_pretrained(str(INTENT_DIR)),
            tokenizer=AutoTokenizer.from_pretrained(str(INTENT_DIR)),
            device=-1, top_k=None,
        )

        print("[chatbot] loading NER model (CPU) ...")
        self.ner_pipe = pipeline(
            "token-classification",
            model=AutoModelForTokenClassification.from_pretrained(str(NER_DIR)),
            tokenizer=AutoTokenizer.from_pretrained(str(NER_DIR)),
            device=-1, aggregation_strategy="simple",
        )

        # RAG (sentence embedder + FAISS) on CPU; embedding 1 query is sub-100ms there.
        print("[chatbot] loading RAG embedder (CPU) ...")
        self.embed_model = SentenceTransformer(EMBED_MODEL, device="cpu")
        self.faiss_index = faiss.read_index(str(RAG_DIR / "faiss.index"))
        self.metadata: list[dict] = json.loads((RAG_DIR / "metadata.json").read_text())

        # Generator on GPU (fp16) if available.
        print(f"[chatbot] loading generator {GEN_MODEL} on {gen_device} (fp16) ...")
        self.gen_tokenizer = AutoTokenizer.from_pretrained(GEN_MODEL)
        self.gen_model = AutoModelForCausalLM.from_pretrained(
            GEN_MODEL,
            torch_dtype=torch.float16 if gpu_available else torch.float32,
            device_map=gen_device,
        )
        self.gen_model.eval()
        self.gen_max_new_tokens = gen_max_new_tokens

        if self.gen_tokenizer.pad_token_id is None:
            self.gen_tokenizer.pad_token_id = self.gen_tokenizer.eos_token_id

        print(f"[chatbot] ready  (KB rows: {len(self.metadata)})")

    # ---- individual stages -------------------------------------------------

    def detect_language(self, text: str) -> tuple[str, float]:
        scores = self.lang_pipe(text)[0]
        best = max(scores, key=lambda s: s["score"])
        return best["label"], float(best["score"])

    def classify_intent(self, text: str) -> tuple[str, float]:
        scores = self.intent_pipe(text)[0]
        best = max(scores, key=lambda s: s["score"])
        return best["label"], float(best["score"])

    def extract_entities(self, text: str) -> list[dict]:
        raw = self.ner_pipe(text)
        out = []
        for ent in raw:
            out.append({
                "text": ent["word"],
                "type": ent["entity_group"],
                "start": int(ent["start"]),
                "end": int(ent["end"]),
                "score": float(ent["score"]),
            })
        return out

    def retrieve(self, query: str, prefer_lang: str | None = None,
                 top_k: int = 3) -> list[dict]:
        q = self.embed_model.encode(
            [query], normalize_embeddings=True, convert_to_numpy=True
        ).astype("float32")
        scores, ids = self.faiss_index.search(q, max(top_k * 2, 6))
        hits = []
        for s, idx in zip(scores[0], ids[0]):
            if idx < 0:
                continue
            row = dict(self.metadata[idx])
            row["score"] = float(s)
            hits.append(row)
        if prefer_lang:
            same = [h for h in hits if h["language"] == prefer_lang]
            other = [h for h in hits if h["language"] != prefer_lang]
            hits = same + other
        return hits[:top_k]

    # ---- generation --------------------------------------------------------

    def _build_messages(self, user_text: str, lang: str,
                        retrieved: list[dict]) -> list[dict]:
        """Build the chat-template messages for Qwen."""
        lang_name = LANG_NAME.get(lang, lang)
        system = (
            "You are a helpful, concise customer-service assistant for a travel and "
            "booking company.\n\n"
            f"REPLY RULES (follow strictly):\n"
            f"- Reply ONLY in {lang_name}. Do not mix languages.\n"
            "- Keep replies to 1–3 sentences.\n"
            "- Write the reply directly. Never start with labels like "
            "'A:', 'Answer:', 'Reply:', or any prefix.\n"
            "- If reference information is provided, use it to inform your "
            "answer but rephrase naturally; do not copy verbatim.\n"
            "- If the references do not cover the question, answer briefly from "
            "general knowledge."
        )
        if retrieved:
            ref_lines = ["Reference information from our knowledge base:"]
            for r in retrieved:
                ref_lines.append(f"- {r['answer']}")
            user = "\n".join(ref_lines) + f"\n\nThe customer says: {user_text}"
        else:
            user = user_text
        return [
            {"role": "system", "content": system},
            {"role": "user", "content": user},
        ]

    @staticmethod
    def _strip_label_prefix(text: str) -> str:
        """Remove stray 'A:'/'Answer:'/etc. prefixes the model might emit."""
        text = text.lstrip()
        for prefix in ("A:", "a:", "Answer:", "Reply:", "Response:",
                       "الإجابة:", "الجواب:", "Réponse:"):
            if text.startswith(prefix):
                return text[len(prefix):].lstrip()
        return text

    def generate_reply(self, user_text: str, lang: str,
                       retrieved: list[dict]) -> str:
        """Synchronous generation (used by CLI / non-streaming callers)."""
        return "".join(self._stream_generate(user_text, lang, retrieved))

    def _stream_generate(self, user_text: str, lang: str,
                         retrieved: list[dict]) -> Iterator[str]:
        """Yield generated text in chunks as it's produced. Threaded internally
        so the generator can drive a UI loop without blocking."""
        messages = self._build_messages(user_text, lang, retrieved)
        prompt = self.gen_tokenizer.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True,
        )
        inputs = self.gen_tokenizer(prompt, return_tensors="pt").to(self.gen_model.device)
        streamer = TextIteratorStreamer(
            self.gen_tokenizer, skip_prompt=True, skip_special_tokens=True,
            timeout=60.0,
        )
        gen_kwargs = dict(
            **inputs,
            max_new_tokens=self.gen_max_new_tokens,
            do_sample=True,
            temperature=0.7,
            top_p=0.9,
            repetition_penalty=1.1,
            pad_token_id=self.gen_tokenizer.pad_token_id,
            streamer=streamer,
        )
        thread = Thread(target=self.gen_model.generate, kwargs=gen_kwargs)
        thread.start()
        try:
            for chunk in streamer:
                yield chunk
        finally:
            thread.join()

    # ---- end-to-end --------------------------------------------------------

    def respond(self, text: str) -> Turn:
        """Blocking response — returns final Turn after full generation."""
        last: Turn | None = None
        for turn in self.respond_stream(text):
            last = turn
        assert last is not None
        return last

    def respond_stream(self, text: str) -> Iterator[Turn]:
        """Streaming response — yields a Turn each time the reply grows.

        For greeting/farewell, yields exactly one Turn with the canned reply.
        For LLM-generated replies, yields once per token chunk so the UI can
        show progressive text.
        """
        text = text.strip()
        if not text:
            yield Turn(text="", language="EN", intent="other",
                       intent_confidence=0.0, entities=[],
                       reply=CANNED["fallback"]["EN"])
            return

        language, _ = self.detect_language(text)
        intent, intent_conf = self.classify_intent(text)
        entities = self.extract_entities(text)

        if intent in ("greeting", "farewell"):
            yield Turn(
                text=text, language=language, intent=intent,
                intent_confidence=intent_conf, entities=entities,
                reply=CANNED[intent].get(language, CANNED[intent]["EN"]),
                retrieved=[], generator_used=False,
            )
            return

        retrieved: list[dict] = []
        if intent in ("inquiry", "booking", "complaint"):
            retrieved = self.retrieve(text, prefer_lang=language, top_k=3)
        # else: "other" — no retrieval, model uses general knowledge.

        running = ""
        for chunk in self._stream_generate(text, language, retrieved):
            running += chunk
            yield Turn(
                text=text, language=language, intent=intent,
                intent_confidence=intent_conf, entities=entities,
                reply=self._strip_label_prefix(running),
                retrieved=retrieved, generator_used=True,
            )
        # Edge case: nothing generated at all.
        if not running.strip():
            yield Turn(
                text=text, language=language, intent=intent,
                intent_confidence=intent_conf, entities=entities,
                reply=CANNED["fallback"].get(language, CANNED["fallback"]["EN"]),
                retrieved=retrieved, generator_used=True,
            )


def _print_turn(turn: Turn) -> None:
    print(f"  language : {turn.language}")
    print(f"  intent   : {turn.intent}  ({turn.intent_confidence:.4f})")
    if turn.entities:
        print("  entities :")
        for e in turn.entities:
            print(f"    [{e['type']}] {e['text']!r}  (score={e['score']:.3f})")
    else:
        print("  entities : (none)")
    if turn.retrieved:
        print(f"  retrieved (top {len(turn.retrieved)}):")
        for r in turn.retrieved:
            print(f"    {r['score']:.3f}  [{r['language']}/{r['topic']}]  {r['question']}")
    print(f"  reply    : {turn.reply}")
    print(f"  source   : {'LLM (Qwen2.5-0.5B)' if turn.generator_used else 'canned'}")


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    parser.add_argument("text", nargs="?", default=None,
                        help="Single user message. If omitted, use --interactive.")
    parser.add_argument("--interactive", action="store_true",
                        help="REPL mode: read messages from stdin until Ctrl-D.")
    parser.add_argument("--json", action="store_true",
                        help="Print one JSON object per turn instead of pretty text.")
    args = parser.parse_args()

    if not args.text and not args.interactive:
        parser.error("provide a message or use --interactive")

    bot = Chatbot()

    def handle(msg: str) -> None:
        turn = bot.respond(msg)
        if args.json:
            print(json.dumps(asdict(turn), ensure_ascii=False))
        else:
            print(f"\n> {msg}")
            _print_turn(turn)

    if args.text:
        handle(args.text)
    if args.interactive:
        print("\n[interactive mode — Ctrl-D / Ctrl-C to exit]")
        try:
            while True:
                msg = input("> ").strip()
                if not msg:
                    continue
                handle(msg)
        except (EOFError, KeyboardInterrupt):
            print("\nbye.")
    return 0


if __name__ == "__main__":
    sys.exit(main())