Deploy updated SLM customer-support chatbot

SLM_CService.py

@@ -1,5 +1,5 @@
 # ── SLM_CService.py ───────────────────────────────────────────────────────────
-# Customer-support-only chatbot with strict NSFW blocking + domain guardrails.
+# Customer-support-only chatbot with strict NSFW blocking + proper Reset.
 
 import os
 import re
@@ -16,26 +16,17 @@ from peft import PeftModel
 from langchain.memory import ConversationBufferMemory
 
 # ──────────────────────────────────────────────────────────────────────────────
-# Hub repos
-REPO = "ThomasBasil/bitext-qlora-tinyllama"   # adapter + tokenizer at ROOT
-BASE = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"   # base model
+REPO = "ThomasBasil/bitext-qlora-tinyllama"
+BASE = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
 
-# Generation params (pass at call time)
 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,
 )
 
-# 4-bit NF4 (GPU needed)
 bnb_cfg = BitsAndBytesConfig(
-    load_in_4bit=True,
-    bnb_4bit_quant_type="nf4",
-    bnb_4bit_use_double_quant=True,
-    bnb_4bit_compute_dtype=torch.float16,  # T4/A10G-friendly
+    load_in_4bit=True, bnb_4bit_quant_type="nf4",
+    bnb_4bit_use_double_quant=True, bnb_4bit_compute_dtype=torch.float16,
 )
 
 # ---- Tokenizer & model -------------------------------------------------------
@@ -45,108 +36,79 @@ 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 (we 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 & blocking (strict)
+# Moderation (unchanged from your last good version)
 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",
     "doggy","missionary","cowgirl","69","kamasutra","dominatrix","submissive","spank",
     "sex position","have sex","make love","how to flirt","dominant in bed",
 ]
-
-def _bad_words_ids(tokenizer, terms: List[str]) -> List[List[int]]:
-    ids = set()
+def _bad_words_ids(tok, terms: List[str]) -> List[List[int]]:
+    ids=set()
     for t in terms:
-        for v in (t, " " + t):
-            toks = tokenizer(v, add_special_tokens=False).input_ids
-            if toks:
-                ids.add(tuple(toks))
+        for v in (t, " "+t):
+            toks = tok(v, add_special_tokens=False).input_ids
+            if toks: ids.add(tuple(toks))
     return [list(t) for t in ids]
-
 BAD_WORD_IDS = _bad_words_ids(tokenizer, SEXUAL_TERMS)
 
 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(item["score"] > 0.60 and item["label"].lower() in
-               {"toxic","severe_toxic","obscene","threat","insult","identity_hate"} for item in scores):
+        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 (kept simple)
-memory = ConversationBufferMemory(return_messages=True)
-
-# System prompt = domain guardrails
+# Memory + globals
+memory = ConversationBufferMemory(return_messages=True)  # has .clear() :contentReference[oaicite:2]{index=2}
 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. "
+    "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 support-ish keywords for routing
 ALLOWED_KEYWORDS = (
     "order","track","status","delivery","shipping","ship","eta","arrive",
     "refund","return","exchange","warranty","guarantee","account","billing",
     "address","cancel","policy","help","support","agent","human"
 )
 
-# FSM helpers
 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
+def extract_order(text: str):
+    m = order_re.search(text); return m.group(1) if m else None
+
 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}"
 def handle_track(o):  return f"Track order #{o} here: https://track.example.com/{o}"
@@ -154,16 +116,31 @@ 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_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?"
 
+# >>> state that must reset <<<
 stored_order   = None
 pending_intent = None
 
+# public reset hook (called from app.py)
+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
+
 # ---- chat templating ---------------------------------------------------------
 def _lc_to_messages() -> List[Dict[str,str]]:
     msgs = [{"role": "system", "content": SYSTEM_PROMPT}]
-    for m in memory.load_memory_variables({}).get("chat_history", []):
+    hist = memory.load_memory_variables({}).get("chat_history", []) or []
+    for m in hist:
         role = "user" if getattr(m, "type", "") == "human" else "assistant"
         msgs.append({"role": role, "content": getattr(m, "content", "")})
     return msgs
@@ -187,7 +164,13 @@ def chat_with_memory(user_input: str) -> str:
     if not ui:
         return "How can I help with your order today?"
 
-    # 1) Safety first
+    # If memory is empty, start clean (fresh session)
+    hist = memory.load_memory_variables({}).get("chat_history", []) or []
+    if len(hist) == 0:
+        stored_order = None
+        pending_intent = None
+
+    # 1) Safety
     if is_sexual_or_toxic(ui):
         reply = REFUSAL
         memory.save_context({"input": ui}, {"output": reply})
@@ -195,7 +178,7 @@ def chat_with_memory(user_input: str) -> str:
 
     low = ui.lower()
 
-    # 2) Quick intents that don't depend on domain keywords
+    # 2) Quick intents
     if any(tok in low for tok in ["thank you","thanks","thx"]):
         reply = handle_gratitude()
         memory.save_context({"input": ui}, {"output": reply})
@@ -205,26 +188,24 @@ def chat_with_memory(user_input: str) -> str:
         memory.save_context({"input": ui}, {"output": reply})
         return reply
 
-    # 3) *** ORDER NUMBER FIRST *** (so follow-ups like "It's #4567" work)
+    # 3) Order number FIRST
     new_o = extract_order(ui)
     if new_o:
         stored_order = new_o
-        if pending_intent in ("status","eta","track","link"):
-            fn = {"status": handle_status,"eta": handle_eta,"track": handle_track,"link": handle_link}[pending_intent]
-            reply = fn(stored_order)
-            pending_intent = None
-            memory.save_context({"input": ui}, {"output": reply})
-            return reply
-        # No pending intent → fall through to classify what they want next.
+        if pending_intent in ("status","eta","track","link","cancel"):
+            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}); return reply
 
-    # 4) Support-only guard (but SKIP if we just saw an order number or have a pending intent)
+    # 4) Support-only guard (skip if pending intent or new order number)
     if pending_intent is None and new_o 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) Intent classification
+    # 5) Intents (added 'cancel')
     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"]):
@@ -233,23 +214,24 @@ 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:
+        intent = "cancel"
     else:
         intent = "fallback"
 
-    # 6) Handle core intents
-    if intent in ("status","eta","track","link"):
+    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}[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
 
-    # 7) LLM fallback (still on-topic) + post-check for safety
+    # 6) LLM fallback (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

app.py

@@ -2,27 +2,30 @@ import os
 os.environ["OMP_NUM_THREADS"] = "1"
 
 import gradio as gr
-from SLM_CService import chat_with_memory
+from SLM_CService import chat_with_memory, reset_state
 
 def respond(user_message, history):
     reply = chat_with_memory(user_message or "")
-    # messages format: list of dicts with role/content
     history = (history or []) + [
         {"role":"user","content": user_message or ""},
         {"role":"assistant","content": reply},
     ]
     return history
 
+def reset_chat():
+    reset_state()          # <-- clear memory + globals
+    return []              # clear UI
+
 with gr.Blocks() as demo:
     gr.Markdown("# 🛎 Customer Support Chatbot")
-    chatbot = gr.Chatbot(type="messages")  # Gradio 'messages' format. :contentReference[oaicite:14]{index=14}
+    chatbot = gr.Chatbot(type="messages")
     with gr.Row():
         user_in = gr.Textbox(placeholder="Ask about orders, tracking, returns…", scale=5)
         send    = gr.Button("Send", variant="primary")
         reset   = gr.Button("🔄 Reset Chat")
     send.click(respond, [user_in, chatbot], [chatbot])
     user_in.submit(respond, [user_in, chatbot], [chatbot])
-    reset.click(lambda: [], None, [chatbot])
+    reset.click(reset_chat, None, [chatbot])   # <-- real reset
 
 if __name__ == "__main__":
     demo.launch(server_name="0.0.0.0", server_port=7860)