Update app.py

app.py

@@ -1,10 +1,11 @@
 """
-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.
+SchoolSpirit AI – Granite‑3.3‑2B chatbot Space
+----------------------------------------------
+• Uses IBM Granite‑3.3‑2B‑Instruct (public, no access token).
+• Fits HF CPU Space (2‑B params, bfloat16).
+• Keeps last MAX_TURNS exchanges.
+• “Clear chat” button resets context.
+• Robust error handling & logging.
 """
 
 import re
@@ -20,23 +21,25 @@ from transformers.utils import logging as hf_logging
 hf_logging.set_verbosity_error()
 LOG = hf_logging.get_logger("SchoolSpirit")
 
-MODEL_ID     = "meta-llama/Llama-3.2-3B-Instruct"
+MODEL_ID     = "ibm-granite/granite-3.3-2b-instruct"
 MAX_TURNS    = 6
-MAX_TOKENS   = 220
-MAX_INPUT_CH = 500
+MAX_TOKENS   = 200
+MAX_INPUT_CH = 400
 
 SYSTEM_MSG = (
-    "You are SchoolSpirit AI, the friendly digital mascot for a company that "
+    "You are SchoolSpirit AI, the upbeat 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, say so and suggest contacting a human. Never request personal data."
+    "hardware for schools. Answer concisely and age‑appropriately. If unsure, "
+    "say so and suggest contacting a human. Do not ask for personal data."
 )
 
 # ────────────────────────── Model Load ──────────────────────────────────────
 try:
     tok = AutoTokenizer.from_pretrained(MODEL_ID)
     model = AutoModelForCausalLM.from_pretrained(
-        MODEL_ID, device_map="auto", torch_dtype="auto"
+        MODEL_ID,
+        device_map="auto",
+        torch_dtype="auto",          # bfloat16/float16 under the hood
     )
     generator = pipeline(
         "text-generation",
@@ -46,81 +49,72 @@ try:
         do_sample=True,
         temperature=0.7,
     )
-    MODEL_LOAD_ERR = None
+    MODEL_ERR = None
 except Exception as exc:  # noqa: BLE001
-    MODEL_LOAD_ERR = f"Model load error: {exc}"
+    MODEL_ERR = f"Model load error: {exc}"
     generator = None
-    LOG.error(MODEL_LOAD_ERR)
+    LOG.error(MODEL_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 truncate(hist):
+    """Return last MAX_TURNS (u,a) 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 "…"
+def clean(text: str) -> str:
+    """Collapse whitespace; never return empty string."""
+    out = re.sub(r"\s+", " ", text.strip())
+    return out or "…"
 
 # ────────────────────────── Chat Callback ───────────────────────────────────
 def chat(history, user_msg):
-    # Gradio guarantees history is a list of tuples (user, bot)
-    history = list(history)
+    history = list(history)  # Gradio ensures list of tuples
 
-    # Start‑up failure fallback
-    if MODEL_LOAD_ERR:
-        history.append((user_msg, MODEL_LOAD_ERR))
+    if MODEL_ERR:
+        history.append((user_msg, MODEL_ERR))
         return history, ""
 
-    # Basic guards
-    user_msg = (user_msg or "").strip()
+    user_msg = clean(user_msg or "")
     if not user_msg:
         history.append(("", "Please enter a message."))
         return history, ""
     if len(user_msg) > MAX_INPUT_CH:
-        history.append(
-            (user_msg, "Sorry, that message is too long. Please shorten it.")
-        )
+        history.append((user_msg, "That message is too long."))
         return history, ""
 
-    history = truncate_history(history, MAX_TURNS)
+    history = truncate(history)
 
     # Build prompt
-    prompt_parts = [SYSTEM_MSG]
+    prompt_lines = [SYSTEM_MSG]
     for u, a in history:
-        prompt_parts += [f"User: {u}", f"AI: {a}"]
-    prompt_parts += [f"User: {user_msg}", "AI:"]
-    prompt = "\n".join(prompt_parts)
+        prompt_lines += [f"User: {u}", f"AI: {a}"]
+    prompt_lines += [f"User: {user_msg}", "AI:"]
+    prompt = "\n".join(prompt_lines)
 
-    # Generate
     try:
-        completion = generator(prompt, truncate=2048)[0]["generated_text"]
-        reply = safe_reply(completion.split("AI:", 1)[-1])
+        completion = generator(prompt, truncate=4096)[0]["generated_text"]
+        reply = clean(completion.split("AI:", 1)[-1])
     except Exception as err:  # noqa: BLE001
         LOG.error(f"Inference error: {err}")
-        reply = (
-            "Sorry—I'm having trouble right now. Please try again in a moment."
-        )
+        reply = "Sorry—I'm having trouble right now. Please try again shortly."
 
     history.append((user_msg, reply))
     return history, ""
 
-# ────────────────────────── Clear Chat fn ───────────────────────────────────
+# ────────────────────────── Clear Chat ──────────────────────────────────────
 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…")
+    chatbot  = gr.Chatbot()
+    msg_box  = 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])
+    send_btn.click(chat, [chatbot, msg_box], [chatbot, msg_box])
+    msg_box.submit(chat, [chatbot, msg_box], [chatbot, msg_box])
+    clear_btn.click(clear_chat, outputs=[chatbot, msg_box])
 
 demo.launch()