Update app.py

app.py

@@ -1,9 +1,10 @@
 """
-SchoolSpirit AI – public chatbot Space
---------------------------------------
-• Tiny Llama‑3 3 B model (fits HF CPU Space).
-• Light‑blue Gradio chat widget.
-• Robust error handling.
+SchoolSpirit AI – hardened chatbot Space with “Clear chat” button
+----------------------------------------------------------------
+• Loads Meta Llama‑3 3 B‑Instruct via transformers.
+• Keeps last MAX_TURNS exchanges to fit context.
+• Adds a “Clear chat” button that resets history.
+• Extensive try/except blocks to prevent widget crashes.
 """
 
 import re
@@ -15,96 +16,111 @@ from transformers import (
 )
 from transformers.utils import logging as hf_logging
 
+# ────────────────────────── Config ──────────────────────────────────────────
 hf_logging.set_verbosity_error()
-LOGGER = hf_logging.get_logger("SchoolSpirit")
+LOG = hf_logging.get_logger("SchoolSpirit")
 
-# ------------------------ Config -------------------------------------------
 MODEL_ID     = "meta-llama/Llama-3.2-3B-Instruct"
-MAX_TURNS    = 6           # last N exchanges kept
-MAX_TOKENS   = 220         # response length
-MAX_INPUT_CH = 500         # user message length guard
+MAX_TURNS    = 6
+MAX_TOKENS   = 220
+MAX_INPUT_CH = 500
 
 SYSTEM_MSG = (
     "You are SchoolSpirit AI, the friendly digital mascot for a company that "
     "offers on‑prem AI chat mascots, fine‑tuning services, and turnkey GPU "
     "hardware for schools. Keep answers concise, upbeat, and age‑appropriate. "
-    "If unsure, admit it and suggest contacting a human. Never request "
-    "personal data."
+    "If unsure, say so and suggest contacting a human. Never request personal data."
 )
 
-# ------------------------ Model Load ---------------------------------------
+# ────────────────────────── Model Load ──────────────────────────────────────
 try:
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
-    model     = AutoModelForCausalLM.from_pretrained(
-        MODEL_ID,
-        device_map="auto",
-        torch_dtype="auto",
+    tok = AutoTokenizer.from_pretrained(MODEL_ID)
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_ID, device_map="auto", torch_dtype="auto"
     )
     generator = pipeline(
         "text-generation",
         model=model,
-        tokenizer=tokenizer,
+        tokenizer=tok,
         max_new_tokens=MAX_TOKENS,
         do_sample=True,
         temperature=0.7,
     )
     MODEL_LOAD_ERR = None
 except Exception as exc:  # noqa: BLE001
-    generator = None
     MODEL_LOAD_ERR = f"Model load error: {exc}"
-    LOGGER.error(MODEL_LOAD_ERR)
+    generator = None
+    LOG.error(MODEL_LOAD_ERR)
+
+# ────────────────────────── Helpers ────────────────────────────────────────
+def truncate_history(hist, max_turns):
+    """Return only the last `max_turns` (user,bot) pairs."""
+    return hist[-max_turns:] if len(hist) > max_turns else hist
+
+
+def safe_reply(msg: str) -> str:
+    """Post‑process model output and ensure it is non‑empty."""
+    msg = msg.strip()
+    msg = re.sub(r"\s+", " ", msg)
+    return msg or "…"
 
-# ------------------------ Chat Function ------------------------------------
+# ────────────────────────── Chat Callback ───────────────────────────────────
 def chat(history, user_msg):
-    """Gradio ChatInterface callback using (history, user_msg) tuples."""
+    # Gradio guarantees history is a list of tuples (user, bot)
+    history = list(history)
 
-    # Hard failure at Space startup
+    # Start‑up failure fallback
     if MODEL_LOAD_ERR:
         history.append((user_msg, MODEL_LOAD_ERR))
         return history, ""
 
-    # Basic user‑input guardrails
+    # Basic guards
     user_msg = (user_msg or "").strip()
     if not user_msg:
-        return history + [("", "Please enter a message.")], ""
+        history.append(("", "Please enter a message."))
+        return history, ""
     if len(user_msg) > MAX_INPUT_CH:
         history.append(
-            (user_msg, "Sorry, your message is too long. Please shorten it.")
+            (user_msg, "Sorry, that message is too long. Please shorten it.")
         )
         return history, ""
 
-    # Keep only last MAX_TURNS
-    if len(history) > MAX_TURNS:
-        history = history[-MAX_TURNS:]
+    history = truncate_history(history, MAX_TURNS)
 
     # Build prompt
-    prompt = [SYSTEM_MSG]
+    prompt_parts = [SYSTEM_MSG]
     for u, a in history:
-        prompt.append(f"User: {u}")
-        prompt.append(f"AI: {a}")
-    prompt.append(f"User: {user_msg}")
-    prompt.append("AI:")
-
-    prompt = "\n".join(prompt)
+        prompt_parts += [f"User: {u}", f"AI: {a}"]
+    prompt_parts += [f"User: {user_msg}", "AI:"]
+    prompt = "\n".join(prompt_parts)
 
-    # Generate reply
+    # Generate
     try:
         completion = generator(prompt, truncate=2048)[0]["generated_text"]
-        reply = completion.split("AI:", 1)[-1].strip()
-        reply = re.sub(r"\s+", " ", reply)  # collapse excess whitespace
+        reply = safe_reply(completion.split("AI:", 1)[-1])
     except Exception as err:  # noqa: BLE001
-        LOGGER.error(f"Inference error: {err}")
+        LOG.error(f"Inference error: {err}")
         reply = (
-            "Sorry, something went wrong on my end. "
-            "Please try again in a few seconds."
+            "Sorry—I'm having trouble right now. Please try again in a moment."
         )
 
     history.append((user_msg, reply))
     return history, ""
 
-# ------------------------ UI -----------------------------------------------
-gr.ChatInterface(
-    chat,
-    title="SchoolSpirit AI Chat",
-    theme=gr.themes.Soft(primary_hue="blue"),  # light‑blue look
-).launch()
+# ────────────────────────── Clear Chat fn ───────────────────────────────────
+def clear_chat():
+    return [], ""
+
+# ────────────────────────── UI Launch ───────────────────────────────────────
+with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue")) as demo:
+    gr.Markdown("# SchoolSpirit AI Chat")
+    chatbot = gr.Chatbot()
+    msg_in  = gr.Textbox(placeholder="Ask me anything about SchoolSpirit AI…")
+    send_btn = gr.Button("Send")
+    clear_btn = gr.Button("Clear Chat", variant="secondary")
+
+    send_btn.click(chat, [chatbot, msg_in], [chatbot, msg_in])
+    msg_in.submit(chat, [chatbot, msg_in], [chatbot, msg_in])
+    clear_btn.click(clear_chat, outputs=[chatbot, msg_in])
+
+demo.launch()