Deploy updated SLM customer-support chatbot

SLM_CService.py

@@ -1,5 +1,5 @@
 # ── SLM_CService.py ───────────────────────────────────────────────────────────
-# Customer-support-only chatbot with strict NSFW blocking + proper Reset.
+# Customer-support-only chatbot with strict NSFW blocking + robust FSM.
 
 import os
 import re
@@ -20,13 +20,19 @@ REPO = "ThomasBasil/bitext-qlora-tinyllama"
 BASE = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
 
 GEN_KW = dict(
-    max_new_tokens=160, do_sample=True, top_p=0.9, temperature=0.7,
-    repetition_penalty=1.1, no_repeat_ngram_size=4,
+    max_new_tokens=160,
+    do_sample=True,
+    top_p=0.9,
+    temperature=0.7,
+    repetition_penalty=1.1,
+    no_repeat_ngram_size=4,
 )
 
 bnb_cfg = BitsAndBytesConfig(
-    load_in_4bit=True, bnb_4bit_quant_type="nf4",
-    bnb_4bit_use_double_quant=True, bnb_4bit_compute_dtype=torch.float16,
+    load_in_4bit=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_use_double_quant=True,
+    bnb_4bit_compute_dtype=torch.float16,
 )
 
 # ---- Tokenizer & model -------------------------------------------------------
@@ -36,22 +42,33 @@ if tokenizer.pad_token_id is None and tokenizer.eos_token_id is not None:
 tokenizer.padding_side = "left"
 tokenizer.truncation_side = "right"
 
+# Unsloth returns (model, tokenizer) → unpack
 model, _ = unsloth.FastLanguageModel.from_pretrained(
-    model_name=BASE, load_in_4bit=True, quantization_config=bnb_cfg,
-    device_map="auto", trust_remote_code=True,
+    model_name=BASE,
+    load_in_4bit=True,
+    quantization_config=bnb_cfg,
+    device_map="auto",
+    trust_remote_code=True,
 )
 unsloth.FastLanguageModel.for_inference(model)
+
+# Apply your PEFT adapter from repo root
 model = PeftModel.from_pretrained(model, REPO)
 model.eval()
 
+# Text-generation pipeline (pass gen params at call time)
 chat_pipe = pipeline(
-    "text-generation", model=model, tokenizer=tokenizer,
-    trust_remote_code=True, return_full_text=False,
+    "text-generation",
+    model=model,
+    tokenizer=tokenizer,
+    trust_remote_code=True,
+    return_full_text=False,
 )
 
 # ──────────────────────────────────────────────────────────────────────────────
-# Moderation (unchanged from your last good version)
+# Moderation & blocking (strict)
 from transformers import TextClassificationPipeline
+
 SEXUAL_TERMS = [
     "sex","sexual","porn","nsfw","fetish","kink","bdsm","nude","naked","anal",
     "blowjob","handjob","cum","breast","boobs","vagina","penis","semen","ejaculate",
@@ -87,27 +104,39 @@ def is_sexual_or_toxic(text: str) -> bool:
             return True
     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)  # has .clear() :contentReference[oaicite:2]{index=2}
+memory = ConversationBufferMemory(return_messages=True)
+
 SYSTEM_PROMPT = (
     "You are a customer-support assistant for our store. Only handle account, "
     "orders, shipping, delivery ETA, tracking links, returns/refunds, warranty, and store policy. "
     "If a request is out of scope or sexual/NSFW, refuse briefly and offer support options. "
     "Be concise and professional."
 )
+
 ALLOWED_KEYWORDS = (
     "order","track","status","delivery","shipping","ship","eta","arrive",
-    "refund","return","exchange","warranty","guarantee","account","billing",
-    "address","cancel","policy","help","support","agent","human"
+    "refund","return","exchange","warranty","guarantee","policy","account","billing",
+    "address","cancel","help","support","agent","human"
+)
+
+# 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,
 )
 
-order_re = re.compile(r"#(\d{1,10})")
 def extract_order(text: str):
-    m = order_re.search(text); return m.group(1) if m else None
+    m = ORDER_RX.search(text or "")
+    if not m: return None
+    return m.group(1) or m.group(2)
 
 def handle_status(o): return f"Order #{o} is in transit and should arrive in 3–5 business days."
 def handle_eta(o):    return f"Delivery for order #{o} typically takes 3–5 days; you can track it at https://track.example.com/{o}"
@@ -116,13 +145,17 @@ def handle_link(o):   return f"Here’s the latest tracking link for order #{o}:
 def handle_return_policy(_=None):
     return ("Our return policy allows returns of unused items in original packaging within 30 days of receipt. "
             "Would you like me to connect you with a human agent?")
+def handle_warranty_policy(_=None):
+    return ("We provide a 1-year limited warranty against manufacturing defects. "
+            "For issues within 30 days, you can return or exchange; after that, warranty service applies. "
+            "Need help starting a claim?")
 def handle_cancel(o=None):
     return (f"I’ve submitted a cancellation request for order #{o}. If it has already shipped, "
             "we’ll process a return/refund once it’s back. You’ll receive a confirmation email shortly.")
 def handle_gratitude(_=None): return "You’re welcome! Anything else I can help with?"
 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
 
@@ -131,7 +164,6 @@ def reset_state():
     global stored_order, pending_intent
     stored_order = None
     pending_intent = None
-    # clear conversation buffer (official API) :contentReference[oaicite:3]{index=3}
     try: memory.clear()
     except Exception: pass
     return True
@@ -164,7 +196,7 @@ def chat_with_memory(user_input: str) -> str:
     if not ui:
         return "How can I help with your order today?"
 
-    # If memory is empty, start clean (fresh session)
+    # Fresh session guard
     hist = memory.load_memory_variables({}).get("chat_history", []) or []
     if len(hist) == 0:
         stored_order = None
@@ -183,12 +215,8 @@ def chat_with_memory(user_input: str) -> str:
         reply = handle_gratitude()
         memory.save_context({"input": ui}, {"output": reply})
         return reply
-    if "return" in low:
-        reply = handle_return_policy()
-        memory.save_context({"input": ui}, {"output": reply})
-        return reply
 
-    # 3) Order number FIRST
+    # 3) Order number FIRST (so “It’s # 67890” completes the prior request)
     new_o = extract_order(ui)
     if new_o:
         stored_order = new_o
@@ -197,15 +225,19 @@ def chat_with_memory(user_input: str) -> str:
                   "link": handle_link,"cancel": handle_cancel}[pending_intent]
             reply = fn(stored_order); pending_intent = None
             memory.save_context({"input": ui}, {"output": reply}); return reply
+        # no pending intent → ask what they want to do with this order
+        reply = handle_ask_action(stored_order)
+        memory.save_context({"input": ui}, {"output": reply})
+        return reply
 
-    # 4) Support-only guard (skip if pending intent or new order number)
-    if pending_intent is None and new_o is None:
+    # 4) Support-only guard (skip if we have a pending intent)
+    if pending_intent is None:
         if not any(k in low for k in ALLOWED_KEYWORDS) and not any(k in low for k in ("hi","hello","hey")):
             reply = "I’m for store support only (orders, shipping, returns, warranty, account). How can I help with those?"
             memory.save_context({"input": ui}, {"output": reply})
             return reply
 
-    # 5) Intents (added 'cancel')
+    # 5) Intent classification (added warranty/guarantee/policy)
     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"]):
@@ -214,11 +246,16 @@ def chat_with_memory(user_input: str) -> str:
         intent = "track"
     elif "tracking link" in low or "resend" in low or "link" in low:
         intent = "link"
-    elif "cancel" in low:
+    elif any(k in low for k in ["cancel","cancellation","abort order"]):
         intent = "cancel"
+    elif any(k in low for k in ["warranty","guarantee","policy"]):
+        intent = "warranty_policy"
+    elif "return" in low:
+        intent = "return_policy"
     else:
         intent = "fallback"
 
+    # 6) Handle intents
     if intent in ("status","eta","track","link","cancel"):
         if not stored_order:
             pending_intent = intent
@@ -227,10 +264,17 @@ def chat_with_memory(user_input: str) -> str:
             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
+        memory.save_context({"input": ui}, {"output": reply}); return reply
+
+    if intent == "warranty_policy":
+        reply = handle_warranty_policy()
+        memory.save_context({"input": ui}, {"output": reply}); return reply
+
+    if intent == "return_policy":
+        reply = handle_return_policy()
+        memory.save_context({"input": ui}, {"output": reply}); return reply
 
-    # 6) LLM fallback (on-topic) + post-check
+    # 7) LLM fallback (on-topic) + post-check
     reply = _generate_reply(ui)
     if is_sexual_or_toxic(reply): reply = REFUSAL
     memory.save_context({"input": ui}, {"output": reply})