Update app.py

app.py

@@ -3,7 +3,9 @@ import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 from transformers.utils import logging as hf_logging
 
-# ---------- Logging ---------------------------------------------------------
+# -------------------------------------------------------------------
+# 1.  Logging helpers
+# -------------------------------------------------------------------
 os.environ["HF_HOME"] = "/data/.huggingface"
 LOG_FILE = "/data/requests.log"
 
@@ -19,26 +21,34 @@ def log(msg: str):
         pass
 
 
-# ---------- Config ----------------------------------------------------------
+# -------------------------------------------------------------------
+# 2.  Configuration
+# -------------------------------------------------------------------
 MODEL_ID = "ibm-granite/granite-3.3-2b-instruct"
-MAX_TURNS, MAX_TOKENS, MAX_INPUT_CH = 6, 128, 400
+MAX_TURNS, MAX_TOKENS, MAX_INPUT_CH = 4, 64, 300
 
 SYSTEM_MSG = (
-    "You are **SchoolSpirit AI**, the digital mascot for SchoolSpirit AI LLC, "
+    "You are **SchoolSpirit AI**, the digital mascot for SchoolSpirit AI LLC, "
     "founded by Charles Norton in 2025. The company installs on‑prem AI chat "
-    "mascots, offers custom fine‑tuning, and ships turnkey GPU hardware to schools.\n\n"
-    "Guidelines:\n"
-    "• Warm, concise answers (max 4 sentences).\n"
+    "mascots, offers custom fine‑tuning, and ships turnkey GPU hardware to "
+    "K‑12 schools.\n\n"
+    "GUIDELINES:\n"
+    "• Warm, encouraging tone for students, parents, staff.\n"
+    "• Replies ≤ 4 sentences unless asked for detail.\n"
+    "• If unsure/out‑of‑scope: say so and suggest human follow‑up.\n"
     "• No personal‑data collection or sensitive advice.\n"
-    "• If unsure, say so and suggest a human follow‑up.\n"
-    "• Avoid profanity, politics, or mature themes."
+    "• No profanity, politics, or mature themes."
 )
-WELCOME_MSG = "Welcome to SchoolSpirit AI! Do you have any questions?"
+WELCOME_MSG = "Welcome to SchoolSpirit AI! Do you have any questions?"
 
-strip = lambda s: re.sub(r"\s+", " ", s.strip())
 
+def strip(s: str) -> str:
+    return re.sub(r"\s+", " ", s.strip())
 
-# ---------- Load model (GPU FP‑16 → CPU fallback) ---------------------------
+
+# -------------------------------------------------------------------
+# 3.  Load model (GPU FP‑16 → CPU fallback)
+# -------------------------------------------------------------------
 hf_logging.set_verbosity_error()
 try:
     log("Loading tokenizer …")
@@ -62,7 +72,6 @@ try:
         max_new_tokens=MAX_TOKENS,
         do_sample=True,
         temperature=0.6,
-        pad_token_id=tok.eos_token_id,
     )
     MODEL_ERR = None
     log("Model loaded ✔")
@@ -71,30 +80,41 @@ except Exception as exc:  # noqa: BLE001
     log(MODEL_ERR)
 
 
-# ---------- Chat callback ---------------------------------------------------
-def chat_fn(user_msg: str, history: list[dict]):
+# -------------------------------------------------------------------
+# 4.  Chat callback
+# -------------------------------------------------------------------
+def chat_fn(user_msg: str, history: list[tuple[str, str]], state: dict):
     """
-    history comes in/out as list[{'role':'user'|'assistant','content':str}, …]
+    history: list of (user, assistant) tuples (Gradio default)
+    state  : dict carrying system_prompt + raw_history for the model
+    Returns updated history (for UI) and state (for next round)
     """
     if MODEL_ERR:
-        return history + [{"role": "assistant", "content": MODEL_ERR}]
+        return history + [(user_msg, MODEL_ERR)], state
 
     user_msg = strip(user_msg or "")
     if not user_msg:
-        return history + [{"role": "assistant", "content": "Please type something."}]
+        return history + [(user_msg, "Please type something.")], state
     if len(user_msg) > MAX_INPUT_CH:
         warn = f"Message too long (>{MAX_INPUT_CH} chars)."
-        return history + [{"role": "assistant", "content": warn}]
-
-    # Append user to history
-    history.append({"role": "user", "content": user_msg})
-
-    # Keep system + last N messages
-    convo = [m for m in history if m["role"] != "system"][-MAX_TURNS * 2 :]
-    prompt_parts = [SYSTEM_MSG] + [
-        f"{'User' if m['role']=='user' else 'AI'}: {m['content']}" for m in convo
-    ] + ["AI:"]
-    prompt = "\n".join(prompt_parts)
+        return history + [(user_msg, warn)], state
+
+    # ------------------------------------------------ Prompt assembly
+    raw_hist = state.get("raw", [])
+    raw_hist.append({"role": "user", "content": user_msg})
+    # keep system + last N exchanges
+    convo = [m for m in raw_hist if m["role"] != "system"][-MAX_TURNS * 2 :]
+    raw_hist = [{"role": "system", "content": SYSTEM_MSG}] + convo
+
+    prompt = "\n".join(
+        [
+            m["content"]
+            if m["role"] == "system"
+            else f'{"User" if m["role"]=="user" else "AI"}: {m["content"]}'
+            for m in raw_hist
+        ]
+        + ["AI:"]
+    )
 
     try:
         raw = gen(prompt)[0]["generated_text"]
@@ -104,22 +124,28 @@ def chat_fn(user_msg: str, history: list[dict]):
         log("❌ Inference error:\n" + traceback.format_exc())
         reply = "Sorry—backend crashed. Please try again later."
 
-    history.append({"role": "assistant", "content": reply})
-    return history
-
-
-# ---------- Launch ----------------------------------------------------------
-gr.ChatInterface(
-    fn=chat_fn,
-    chatbot=gr.Chatbot(
-        height=480,
-        type="messages",
-        value=[
-            {"role": "assistant", "content": WELCOME_MSG}
-        ],  # ONE welcome bubble
-    ),
-    additional_inputs=None,
-    title="SchoolSpirit AI Chat",
-    theme=gr.themes.Soft(primary_hue="blue"),
-    examples=None,
-).launch()
+    # ------------------------------------------------ Update state + UI history
+    raw_hist.append({"role": "assistant", "content": reply})
+    state["raw"] = raw_hist
+    history.append((user_msg, reply))
+    return history, state
+
+
+# -------------------------------------------------------------------
+# 5.  Launch
+# -------------------------------------------------------------------
+with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue")) as demo:
+    chatbot = gr.Chatbot(
+        value=[("", WELCOME_MSG)], height=480, label="SchoolSpirit AI"
+    )
+    state = gr.State({"raw": [{"role": "system", "content": SYSTEM_MSG}]})
+    with gr.Row():
+        txt = gr.Textbox(
+            scale=4, placeholder="Type your question here...", show_label=False
+        )
+        send = gr.Button("Send", variant="primary")
+
+    send.click(chat_fn, inputs=[txt, chatbot, state], outputs=[chatbot, state])
+    txt.submit(chat_fn, inputs=[txt, chatbot, state], outputs=[chatbot, state])
+
+demo.launch()