Update app.py

app.py

@@ -1,62 +1,110 @@
+"""
+SchoolSpirit AI – public chatbot Space
+--------------------------------------
+• Tiny Llama‑3 3 B model (fits HF CPU Space).
+• Light‑blue Gradio chat widget.
+• Robust error handling.
+"""
+
+import re
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+from transformers import (
+    AutoTokenizer,
+    AutoModelForCausalLM,
+    pipeline,
+)
 from transformers.utils import logging as hf_logging
 
 hf_logging.set_verbosity_error()
+LOGGER = 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
 
-MODEL_ID   = "meta-llama/Llama-3.2-3B-Instruct"
-MAX_TURNS  = 6
-MAX_TOKENS = 220
 SYSTEM_MSG = (
     "You are SchoolSpirit AI, the friendly digital mascot for a company that "
-    "provides on‑prem AI chat mascots, fine‑tuning services, and turnkey GPU "
-    "hardware for schools.  Keep answers concise, upbeat, and age‑appropriate.  "
-    "If you don’t know, say so and suggest contacting a human.  Never request "
+    "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."
 )
 
+# ------------------------ Model Load ---------------------------------------
 try:
-    tok = AutoTokenizer.from_pretrained(MODEL_ID)
-    model = AutoModelForCausalLM.from_pretrained(
-        MODEL_ID, device_map="auto", torch_dtype="auto"
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+    model     = AutoModelForCausalLM.from_pretrained(
+        MODEL_ID,
+        device_map="auto",
+        torch_dtype="auto",
     )
-    gen = pipeline(
+    generator = pipeline(
         "text-generation",
         model=model,
-        tokenizer=tok,
+        tokenizer=tokenizer,
         max_new_tokens=MAX_TOKENS,
         do_sample=True,
         temperature=0.7,
     )
-    model_error = None
+    MODEL_LOAD_ERR = None
 except Exception as exc:  # noqa: BLE001
-    model_error = f"Model load error: {exc}"
-    gen = None
+    generator = None
+    MODEL_LOAD_ERR = f"Model load error: {exc}"
+    LOGGER.error(MODEL_LOAD_ERR)
 
+# ------------------------ Chat Function ------------------------------------
 def chat(history, user_msg):
-    if model_error:
-        return history + [(user_msg, model_error)], ""
+    """Gradio ChatInterface callback using (history, user_msg) tuples."""
 
+    # Hard failure at Space startup
+    if MODEL_LOAD_ERR:
+        history.append((user_msg, MODEL_LOAD_ERR))
+        return history, ""
+
+    # Basic user‑input guardrails
+    user_msg = (user_msg or "").strip()
+    if not user_msg:
+        return history + [("", "Please enter a message.")], ""
+    if len(user_msg) > MAX_INPUT_CH:
+        history.append(
+            (user_msg, "Sorry, your message is too long. Please shorten it.")
+        )
+        return history, ""
+
+    # Keep only last MAX_TURNS
     if len(history) > MAX_TURNS:
         history = history[-MAX_TURNS:]
 
-    prompt = SYSTEM_MSG + "\n"
+    # Build prompt
+    prompt = [SYSTEM_MSG]
     for u, a in history:
-        prompt += f"User: {u}\nAI: {a}\n"
-    prompt += f"User: {user_msg}\nAI:"
+        prompt.append(f"User: {u}")
+        prompt.append(f"AI: {a}")
+    prompt.append(f"User: {user_msg}")
+    prompt.append("AI:")
+
+    prompt = "\n".join(prompt)
 
+    # Generate reply
     try:
-        completion = gen(prompt)[0]["generated_text"]
+        completion = generator(prompt, truncate=2048)[0]["generated_text"]
         reply = completion.split("AI:", 1)[-1].strip()
+        reply = re.sub(r"\s+", " ", reply)  # collapse excess whitespace
     except Exception as err:  # noqa: BLE001
-        reply = "Sorry, an internal error occurred. Please try again later."
-        hf_logging.get_logger("SchoolSpirit").error(str(err))
+        LOGGER.error(f"Inference error: {err}")
+        reply = (
+            "Sorry, something went wrong on my end. "
+            "Please try again in a few seconds."
+        )
 
     history.append((user_msg, reply))
     return history, ""
 
+# ------------------------ UI -----------------------------------------------
 gr.ChatInterface(
     chat,
     title="SchoolSpirit AI Chat",
-    theme=gr.themes.Soft(primary_hue="blue"),
+    theme=gr.themes.Soft(primary_hue="blue"),  # light‑blue look
 ).launch()