Deploy updated SLM customer-support chatbot

SLM_CService.py

@@ -1,25 +1,30 @@
 # ── SLM_CService.py ───────────────────────────────────────────────────────────
-# Customer-support-only chatbot with strict NSFW blocking + robust FSM.
+# Customer-support-only chatbot with strict NSFW blocking + robust FSM + proper reset.
 
 import os
 import re
 from typing import List, Dict
 
+# Keep OpenMP logs quiet
 os.environ["OMP_NUM_THREADS"] = "1"
+# Ensure we don't accidentally force offline mode
 os.environ.pop("HF_HUB_OFFLINE", None)
 
-# Unsloth must come before transformers/peft
+# ── Import order matters: Unsloth should come before transformers/peft.
 import unsloth  # noqa: E402
+
 import torch
 from transformers import AutoTokenizer, BitsAndBytesConfig, pipeline
 from peft import PeftModel
 from langchain.memory import ConversationBufferMemory
 
-# ──────────────────────────────────────────────────────────────────────────────
-REPO = "ThomasBasil/bitext-qlora-tinyllama"
-BASE = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
+# ==============================
+# Config
+# ==============================
+REPO = "ThomasBasil/bitext-qlora-tinyllama"   # your adapter + tokenizer live at repo root
+BASE = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"   # base model
 
-GEN_KW = dict(
+GEN_KW = dict(  # generation params (passed at call time)
     max_new_tokens=160,
     do_sample=True,
     top_p=0.9,
@@ -28,21 +33,26 @@ GEN_KW = dict(
     no_repeat_ngram_size=4,
 )
 
-bnb_cfg = BitsAndBytesConfig(
+bnb_cfg = BitsAndBytesConfig(  # 4-bit QLoRA-style loading (needs GPU)
     load_in_4bit=True,
     bnb_4bit_quant_type="nf4",
     bnb_4bit_use_double_quant=True,
-    bnb_4bit_compute_dtype=torch.float16,
+    bnb_4bit_compute_dtype=torch.float16,  # T4/A10G-friendly
 )
 
-# ---- Tokenizer & model -------------------------------------------------------
+# Memory key FIX: use the same key for saving & reading history
+MEMORY_KEY = "chat_history"
+
+# ==============================
+# Load tokenizer & model
+# ==============================
 tokenizer = AutoTokenizer.from_pretrained(REPO, use_fast=False)
 if tokenizer.pad_token_id is None and tokenizer.eos_token_id is not None:
     tokenizer.pad_token_id = tokenizer.eos_token_id
 tokenizer.padding_side = "left"
 tokenizer.truncation_side = "right"
 
-# Unsloth returns (model, tokenizer) → unpack
+# Unsloth returns (model, tokenizer) -> unpack
 model, _ = unsloth.FastLanguageModel.from_pretrained(
     model_name=BASE,
     load_in_4bit=True,
@@ -52,11 +62,11 @@ model, _ = unsloth.FastLanguageModel.from_pretrained(
 )
 unsloth.FastLanguageModel.for_inference(model)
 
-# Apply your PEFT adapter from repo root
+# Attach your PEFT adapter from repo root
 model = PeftModel.from_pretrained(model, REPO)
 model.eval()
 
-# Text-generation pipeline (pass gen params at call time)
+# Text-generation pipeline (pass GEN_KW at call time, not as generate_kwargs)
 chat_pipe = pipeline(
     "text-generation",
     model=model,
@@ -65,52 +75,74 @@ chat_pipe = pipeline(
     return_full_text=False,
 )
 
-# ──────────────────────────────────────────────────────────────────────────────
-# Moderation & blocking (strict)
+# ==============================
+# Moderation (strict)
+# ==============================
 from transformers import TextClassificationPipeline
 
 SEXUAL_TERMS = [
+    # single words
     "sex","sexual","porn","nsfw","fetish","kink","bdsm","nude","naked","anal",
     "blowjob","handjob","cum","breast","boobs","vagina","penis","semen","ejaculate",
     "doggy","missionary","cowgirl","69","kamasutra","dominatrix","submissive","spank",
+    # phrases
     "sex position","have sex","make love","how to flirt","dominant in bed",
 ]
+
 def _bad_words_ids(tok, terms: List[str]) -> List[List[int]]:
-    ids=set()
+    """Build bad_words_ids for generation; include both 'term' and ' term' variants."""
+    ids = set()
     for t in terms:
-        for v in (t, " "+t):
+        for v in (t, " " + t):
             toks = tok(v, add_special_tokens=False).input_ids
-            if toks: ids.add(tuple(toks))
+            if toks:
+                ids.add(tuple(toks))
     return [list(t) for t in ids]
+
 BAD_WORD_IDS = _bad_words_ids(tokenizer, SEXUAL_TERMS)
 
+# Lightweight classifiers (optional but helpful defense-in-depth)
 nsfw_cls: TextClassificationPipeline = pipeline(
-    "text-classification", model="eliasalbouzidi/distilbert-nsfw-text-classifier", truncation=True,
+    "text-classification",
+    model="eliasalbouzidi/distilbert-nsfw-text-classifier",
+    truncation=True,
 )
 toxicity_cls: TextClassificationPipeline = pipeline(
-    "text-classification", model="unitary/toxic-bert", truncation=True, return_all_scores=True,
+    "text-classification",
+    model="unitary/toxic-bert",
+    truncation=True,
+    return_all_scores=True,
 )
+
 def is_sexual_or_toxic(text: str) -> bool:
     t = (text or "").lower()
-    if any(k in t for k in SEXUAL_TERMS): return True
+    if any(k in t for k in SEXUAL_TERMS):
+        return True
     try:
         res = nsfw_cls(t)[0]
-        if (res.get("label","").lower()=="nsfw") and float(res.get("score",0))>0.60: return True
-    except Exception: pass
+        if (res.get("label","").lower() == "nsfw") and float(res.get("score",0)) > 0.60:
+            return True
+    except Exception:
+        pass
     try:
         scores = toxicity_cls(t)[0]
-        if any(s["score"]>0.60 and s["label"].lower() in
+        if any(s["score"] > 0.60 and s["label"].lower() in
                {"toxic","severe_toxic","obscene","threat","insult","identity_hate"} for s in scores):
             return True
-    except Exception: pass
+    except Exception:
+        pass
     return False
 
 REFUSAL = ("Sorry, I can’t help with that. I’m only for store support "
            "(orders, shipping, ETA, tracking, returns, warranty, account).")
 
-# ──────────────────────────────────────────────────────────────────────────────
-# Memory + globals
-memory = ConversationBufferMemory(return_messages=True)
+# ==============================
+# Memory + Globals
+# ==============================
+memory = ConversationBufferMemory(
+    memory_key=MEMORY_KEY,      # ← FIX: explicit memory key
+    return_messages=True,
+)
 
 SYSTEM_PROMPT = (
     "You are a customer-support assistant for our store. Only handle account, "
@@ -126,12 +158,16 @@ ALLOWED_KEYWORDS = (
 )
 
 # Robust order detection:
+# - "#67890" / "# 67890"
+# - "order 67890", "order no. 67890", "order number 67890", "order id 67890"
 ORDER_RX = re.compile(
     r"(?:#\s*([\d]{3,12})|order(?:\s*(?:no\.?|number|id))?\s*#?\s*([\d]{3,12}))",
     flags=re.I,
 )
+
 def extract_order(text: str):
-    if not text: return None
+    if not text:
+        return None
     m = ORDER_RX.search(text)
     return (m.group(1) or m.group(2)) if m else None
 
@@ -153,47 +189,56 @@ def handle_gratitude(_=None): return "You’re welcome! Anything else I can help
 def handle_escalation(_=None): return "I can connect you with a human agent. Would you like me to do that?"
 def handle_ask_action(o): return (f"I’ve saved order #{o}. What would you like to do — status, ETA, tracking link, or cancel?")
 
+# >>> state that must reset <<<
 stored_order   = None
 pending_intent = None
 
 def reset_state():
+    """Called by app.py Reset button to clear memory + globals."""
     global stored_order, pending_intent
     stored_order = None
     pending_intent = None
-    try: memory.clear()
-    except Exception: pass
+    try:
+        memory.clear()  # wipe the buffer
+    except Exception:
+        pass
     return True
 
-# ---- chat templating ---------------------------------------------------------
-def _lc_to_messages() -> List[Dict[str,str]]:
+# ==============================
+# Chat templating helpers
+# ==============================
+def _lc_to_messages() -> List[Dict[str, str]]:
     msgs = [{"role": "system", "content": SYSTEM_PROMPT}]
-    hist = memory.load_memory_variables({}).get("chat_history", []) or []
+    hist = memory.load_memory_variables({}).get(MEMORY_KEY, []) or []  # ← use same key
     for m in hist:
         role = "user" if getattr(m, "type", "") == "human" else "assistant"
         msgs.append({"role": role, "content": getattr(m, "content", "")})
     return msgs
 
 def _generate_reply(user_input: str) -> str:
+    # Format with HF chat template so the model respects roles/system
     messages = _lc_to_messages() + [{"role": "user", "content": user_input}]
     prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
     out = chat_pipe(
         prompt,
         eos_token_id=tokenizer.eos_token_id,
         pad_token_id=tokenizer.pad_token_id,
-        bad_words_ids=BAD_WORD_IDS,
+        bad_words_ids=BAD_WORD_IDS,  # block sexual tokens at generation time
         **GEN_KW,
     )[0]["generated_text"]
     return out.strip()
 
-# ---- main entry --------------------------------------------------------------
+# ==============================
+# Main entry
+# ==============================
 def chat_with_memory(user_input: str) -> str:
     global stored_order, pending_intent
     ui = (user_input or "").strip()
     if not ui:
         return "How can I help with your order today?"
 
-    # Fresh session guard
-    hist = memory.load_memory_variables({}).get("chat_history", []) or []
+    # Fresh session guard: if memory empty, also clear globals
+    hist = memory.load_memory_variables({}).get(MEMORY_KEY, []) or []
     if len(hist) == 0:
         stored_order = None
         pending_intent = None
@@ -206,20 +251,24 @@ def chat_with_memory(user_input: str) -> str:
 
     low = ui.lower()
 
-    # 2) Quick intents
+    # 2) Quick intents (gratitude / returns)
     if any(tok in low for tok in ["thank you","thanks","thx"]):
         reply = handle_gratitude()
         memory.save_context({"input": ui}, {"output": reply})
         return reply
 
-    # 3) PENDING-INTENT SHORT-CIRCUIT (fix for "It's #26790" case)
+    # 3) PENDING-INTENT SHORT-CIRCUIT (fixes "It's #26790" case)
     new_o = extract_order(ui)
     if pending_intent:
         if new_o:
             stored_order = new_o
-            fn = {"status": handle_status, "eta": handle_eta,
-                  "track": handle_track, "link": handle_link,
-                  "cancel": handle_cancel}[pending_intent]
+            fn = {
+                "status": handle_status,
+                "eta":    handle_eta,
+                "track":  handle_track,
+                "link":   handle_link,
+                "cancel": handle_cancel,
+            }[pending_intent]
             reply = fn(stored_order)
             pending_intent = None
             memory.save_context({"input": ui}, {"output": reply})
@@ -229,7 +278,7 @@ def chat_with_memory(user_input: str) -> str:
         memory.save_context({"input": ui}, {"output": reply})
         return reply
 
-    # 4) If message provides an order number (no pending intent yet), save it & ask action
+    # 4) If message provides an order number (no pending intent yet), save & ask action
     if new_o:
         stored_order = new_o
         reply = handle_ask_action(stored_order)
@@ -242,7 +291,7 @@ def chat_with_memory(user_input: str) -> str:
         memory.save_context({"input": ui}, {"output": reply})
         return reply
 
-    # 6) Intent classification
+    # 6) Intent classification (deterministic handlers first)
     if any(k in low for k in ["status","where is my order","check status"]):
         intent = "status"
     elif any(k in low for k in ["how long","eta","delivery time"]):
@@ -260,18 +309,24 @@ def chat_with_memory(user_input: str) -> str:
     else:
         intent = "fallback"
 
-    # 7) Handle intents
+    # 7) Handle intents that need an order number
     if intent in ("status","eta","track","link","cancel"):
         if not stored_order:
             pending_intent = intent
             reply = "Sure—what’s your order number (e.g., #12345)?"
         else:
-            fn = {"status": handle_status,"eta": handle_eta,"track": handle_track,
-                  "link": handle_link,"cancel": handle_cancel}[intent]
+            fn = {
+                "status": handle_status,
+                "eta":    handle_eta,
+                "track":  handle_track,
+                "link":   handle_link,
+                "cancel": handle_cancel,
+            }[intent]
             reply = fn(stored_order)
         memory.save_context({"input": ui}, {"output": reply})
         return reply
 
+    # 8) Policy intents (no order needed)
     if intent == "warranty_policy":
         reply = handle_warranty_policy()
         memory.save_context({"input": ui}, {"output": reply})
@@ -282,8 +337,9 @@ def chat_with_memory(user_input: str) -> str:
         memory.save_context({"input": ui}, {"output": reply})
         return reply
 
-    # 8) LLM fallback (on-topic) + post-check
+    # 9) LLM fallback (still on-topic) + post-check
     reply = _generate_reply(ui)
-    if is_sexual_or_toxic(reply): reply = REFUSAL
+    if is_sexual_or_toxic(reply):
+        reply = REFUSAL
     memory.save_context({"input": ui}, {"output": reply})
     return reply