Update app.py

app.py

@@ -1,21 +1,20 @@
 """
-SchoolSpirit AI – minimal public chatbot Space
----------------------------------------------
-• Loads Meta’s Llama‑3 3 B‑Instruct (fits HF CPU Space).
-• Uses Hugging Face transformers + Gradio; no external deps.
-• Keeps prompt short and trims history to fit the model context.
-• Gracefully handles model‑load or inference errors.
+SchoolSpirit AI – Llama‑3 3 B public chatbot Space
+-------------------------------------------------
+• Loads Meta Llama‑3.2‑3B‑Instruct.
+• Keeps only last 6 turns to fit context.
+• Handles model‑load or generation failures gracefully.
 """
 
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 from transformers.utils import logging as hf_logging
 
-hf_logging.set_verbosity_error()  # keep Space logs clean
+hf_logging.set_verbosity_error()
 
 MODEL_ID   = "meta-llama/Llama-3.2-3B-Instruct"
-MAX_TURNS  = 6          # retain last N exchanges
-MAX_TOKENS = 220        # response length
+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 "
@@ -24,55 +23,71 @@ SYSTEM_MSG = (
     "personal data."
 )
 
-# ---------------- Model ------------------------------------------------------
+# ---------------------------------------------------------------------------
 try:
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
-    model     = AutoModelForCausalLM.from_pretrained(
-        MODEL_ID,
-        device_map="auto",          # auto‑detect CPU / GPU if available
-        torch_dtype="auto"
+    tok = AutoTokenizer.from_pretrained(MODEL_ID)
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_ID, device_map="auto", torch_dtype="auto"
     )
     gen = pipeline(
         "text-generation",
         model=model,
-        tokenizer=tokenizer,
+        tokenizer=tok,
         max_new_tokens=MAX_TOKENS,
         do_sample=True,
         temperature=0.7,
     )
-except Exception as e:  # noqa: BLE001
-    # Fatal startup failure – expose error in UI
-    def chat(history, user_msg):
-        return history + [(user_msg, f"Model load error: {e}")], ""
-else:
-    # ---------------- Chat handler ------------------------------------------
-    def chat(history, user_msg):
-        """Gradio ChatInterface callback."""
-        # Trim history to last N turns
-        if len(history) > MAX_TURNS:
-            history = history[-MAX_TURNS:]
+    model_error = None
+except Exception as exc:  # noqa: BLE001
+    model_error = f"Model load error: {exc}"
+    gen = None  # ensure var exists
+# ---------------------------------------------------------------------------
 
-        # Build prompt
-        prompt = SYSTEM_MSG + "\n"
-        for u, a in history:
-            prompt += f"User: {u}\nAI: {a}\n"
-        prompt += f"User: {user_msg}\nAI:"
 
-        # Generate
-        try:
-            completion = gen(prompt)[0]["generated_text"]
-            reply = completion.split("AI:", 1)[-1].strip()
-        except Exception as err:  # noqa: BLE001
-            reply = "Sorry, something went wrong. Please try again later."
-            hf_logging.get_logger("SchoolSpirit").error(str(err))
+def chat(history, user_msg):
+    """Gradio ChatInterface callback using new 'messages' format."""
+    if model_error:
+        history.append(
+            {
+                "role": "assistant",
+                "content": model_error,
+            }
+        )
+        return history
 
-        history.append((user_msg, reply))
-        return history, ""
+    # Trim to last MAX_TURNS messages (role+assistant pairs)
+    if len(history) > MAX_TURNS * 2:
+        history = history[-MAX_TURNS * 2 :]
+
+    # Build prompt
+    prompt = SYSTEM_MSG + "\n"
+    for msg in history:
+        role = "User" if msg["role"] == "user" else "AI"
+        prompt += f"{role}: {msg['content']}\n"
+    prompt += f"User: {user_msg}\nAI:"
+
+    try:
+        completion = gen(prompt)[0]["generated_text"]
+        reply = completion.split("AI:", 1)[-1].strip()
+    except Exception as err:  # noqa: BLE001
+        reply = (
+            "Sorry, something went wrong on my end. "
+            "Please try again in a few seconds."
+        )
+        hf_logging.get_logger("SchoolSpirit").error(str(err))
+
+    history.extend(
+        [
+            {"role": "user", "content": user_msg},
+            {"role": "assistant", "content": reply},
+        ]
+    )
+    return history
 
 
-# ---------------- UI ---------------------------------------------------------
 gr.ChatInterface(
     chat,
-    title="SchoolSpirit AI Chat",
-    theme=gr.themes.Soft(primary_hue="blue"),  # light‑blue chat UI
+    title="SchoolSpirit AI Chat",
+    theme=gr.themes.Soft(primary_hue="blue"),
+    type="messages",  # avoids “tuples” deprecation warning
 ).launch()