Update app.py

app.py

@@ -1,14 +1,15 @@
-##############################################################################
-# SchoolSpirit AI Chat – full‑context, duplicate‑free implementation
-##############################################################################
-
-import os, re, time, datetime, traceback, torch
+import os
+import re
+import time
+import datetime
+import traceback
+import torch
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 from transformers.utils import logging as hf_logging
 
 # ---------------------------------------------------------------------------
-# 0. Basic logging
+# 0.  Paths & basic logging helper
 # ---------------------------------------------------------------------------
 os.environ["HF_HOME"] = "/data/.huggingface"
 LOG_FILE = "/data/requests.log"
@@ -26,24 +27,24 @@ def log(msg: str):
 
 
 # ---------------------------------------------------------------------------
-# 1. Configuration
+# 1.  Configuration constants
 # ---------------------------------------------------------------------------
-MODEL_ID = "ibm-granite/granite-3.3-2b-instruct"  # 2‑B local model
-CONTEXT_TOKENS = 1800      # keep prompt below this many tokens
+MODEL_ID = "ibm-granite/granite-3.3-2b-instruct"  # 2 B model fits Spaces
+CONTEXT_TOKENS = 1800   # leave head‑room for reply inside 2k window
 MAX_NEW_TOKENS = 64
 TEMPERATURE = 0.6
-MAX_INPUT_CH = 300
+MAX_INPUT_CH = 300      # UI safeguard
 
 SYSTEM_MSG = (
-    "You are **SchoolSpirit AI**, digital mascot of SchoolSpirit AI LLC, "
-    "founded by Charles Norton in 2025. The company deploys on‑prem AI chat "
-    "mascots, fine‑tunes language models, and ships turnkey GPU servers to "
-    "K‑12 schools.\n\n"
+    "You are **SchoolSpirit AI**, the official digital mascot of "
+    "SchoolSpirit AI LLC.  Founded by Charles Norton in 2025, the company "
+    "deploys on‑prem AI chat mascots, fine‑tunes language models, and ships "
+    "turnkey GPU servers to K‑12 schools.\n\n"
     "RULES:\n"
-    "• Respond warmly and concisely (≤4 sentences unless asked for detail).\n"
-    "• No personal‑data collection; no medical/legal/financial advice.\n"
-    "• Admit uncertainty and offer human follow‑up if unsure.\n"
-    "• Avoid profanity, politics, mature themes."
+    "• Friendly, concise (≤4 sentences unless prompted).\n"
+    "• No personal data collection; no medical/legal/financial advice.\n"
+    "• If uncertain, admit it & suggest human follow‑up.\n"
+    "• avoid profanity, politics, mature themes."
 )
 WELCOME_MSG = "Welcome to SchoolSpirit AI! Do you have any questions?"
 
@@ -51,19 +52,20 @@ strip = lambda s: re.sub(r"\s+", " ", s.strip())
 
 
 # ---------------------------------------------------------------------------
-# 2. Load tokenizer + model (GPU FP‑16 → CPU fallback)
+# 2.  Load tokenizer + model  (GPU FP‑16 → CPU)
 # ---------------------------------------------------------------------------
 hf_logging.set_verbosity_error()
 try:
+    log("Loading tokenizer …")
     tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
 
     if torch.cuda.is_available():
-        log("GPU found → loading model in FP‑16")
+        log("GPU detected → FP‑16")
         model = AutoModelForCausalLM.from_pretrained(
             MODEL_ID, device_map="auto", torch_dtype=torch.float16
         )
     else:
-        log("No GPU → CPU load")
+        log("CPU fallback")
         model = AutoModelForCausalLM.from_pretrained(
             MODEL_ID,
             device_map="cpu",
@@ -76,11 +78,12 @@ try:
         model=model,
         tokenizer=tokenizer,
         max_new_tokens=MAX_NEW_TOKENS,
-        temperature=TEMPERATURE,
         do_sample=True,
+        temperature=TEMPERATURE,
+        return_full_text=False,    # ← only return the newly generated text
     )
     MODEL_ERR = None
-    log("Model ready ✔")
+    log("Model loaded ✔")
 except Exception as exc:
     MODEL_ERR = f"Model load error: {exc}"
     generator = None
@@ -88,77 +91,79 @@ except Exception as exc:
 
 
 # ---------------------------------------------------------------------------
-# 3. Prompt builder that respects a token budget
+# 3.  Helper: build prompt under token budget
 # ---------------------------------------------------------------------------
-def build_prompt(history_raw: list[dict]) -> str:
+def build_prompt(raw_history: list[dict]) -> str:
     """
-    Accepts full message list. Drops oldest user+assistant pairs until the
-    prompt’s token length is ≤ CONTEXT_TOKENS.
+    raw_history: list  [{'role':'system'|'user'|'assistant', 'content': str}, ...]
+    Keeps trimming oldest user/assistant pair until total tokens < CONTEXT_TOKENS
     """
     def render(msg):
         if msg["role"] == "system":
-            return f"System: {msg['content']}"
+            return msg["content"]
         prefix = "User:" if msg["role"] == "user" else "AI:"
         return f"{prefix} {msg['content']}"
 
-    # split out system + conversation
-    system = [m for m in history_raw if m["role"] == "system"][0]
-    convo = [m for m in history_raw if m["role"] != "system"]
+    # always include system
+    system_msg = [msg for msg in raw_history if msg["role"] == "system"][0]
+    convo = [m for m in raw_history if m["role"] != "system"]
 
+    # iterative trim
     while True:
-        prompt_parts = [render(system)] + [render(m) for m in convo] + ["AI:"]
-        tokens = len(tokenizer.encode("\n".join(prompt_parts), add_special_tokens=False))
-        if tokens <= CONTEXT_TOKENS or len(convo) <= 2:
+        prompt_parts = [system_msg["content"]] + [render(m) for m in convo] + ["AI:"]
+        token_len = len(tokenizer.encode("\n".join(prompt_parts), add_special_tokens=False))
+        if token_len <= CONTEXT_TOKENS or len(convo) <= 2:
             break
-        # remove oldest user+assistant pair
+        # drop oldest user+assistant pair
         convo = convo[2:]
 
     return "\n".join(prompt_parts)
 
 
 # ---------------------------------------------------------------------------
-# 4. Chat callback
+# 4.  Chat callback
 # ---------------------------------------------------------------------------
-def chat_fn(user_text, chat_ui, state):
+def chat_fn(user_msg: str, display_history: list, state: dict):
     """
-    chat_ui : list[(user, assistant)] displayed by gr.Chatbot
-    state   : dict{ 'history_raw': [...] } used for prompt construction
+    display_history : list[tuple[str,str]] for UI
+    state["raw"]    : list[dict]          for prompting
     """
-    user_text = strip(user_text or "")
-    if not user_text:
-        return chat_ui, state
+    user_msg = strip(user_msg or "")
+    if not user_msg:
+        return display_history, state
 
-    if len(user_text) > MAX_INPUT_CH:
-        chat_ui.append((user_text, f"Input exceeds {MAX_INPUT_CH} characters."))
-        return chat_ui, state
+    if len(user_msg) > MAX_INPUT_CH:
+        display_history.append((user_msg, f"Input >{MAX_INPUT_CH} chars."))
+        return display_history, state
 
     if MODEL_ERR:
-        chat_ui.append((user_text, MODEL_ERR))
-        return chat_ui, state
+        display_history.append((user_msg, MODEL_ERR))
+        return display_history, state
 
-    # update raw history (system already present)
-    state["history_raw"].append({"role": "user", "content": user_text})
+    # --- Update raw history
+    state["raw"].append({"role": "user", "content": user_msg})
 
-    prompt = build_prompt(state["history_raw"])
+    # --- Build prompt within token budget
+    prompt = build_prompt(state["raw"])
 
+    # --- Generate
     try:
         start = time.time()
-        out = generator(prompt)[0]["generated_text"]
-        reply = strip(out.split("AI:", 1)[-1])
-        reply = re.split(r"\b(?:User:|AI:)", reply, 1)[0].strip()
-        log(f"Generated in {time.time()-start:.2f}s ({len(reply)} chars)")
+        result = generator(prompt)[0]
+        reply = strip(result["generated_text"])
+        log(f"Reply in {time.time() - start:.2f}s ({len(reply)} chars)")
     except Exception:
-        log("⚠️  Inference error\n" + traceback.format_exc())
-        reply = "Sorry — an internal error occurred. Please try again."
+        log("❌ Inference error:\n" + traceback.format_exc())
+        reply = "Apologies—an internal error occurred. Please try again."
 
-    # append to histories
-    chat_ui.append((user_text, reply))
-    state["history_raw"].append({"role": "assistant", "content": reply})
-    return chat_ui, state
+    # --- Append assistant reply to both histories
+    display_history.append((user_msg, reply))
+    state["raw"].append({"role": "assistant", "content": reply})
+    return display_history, state
 
 
 # ---------------------------------------------------------------------------
-# 5. Gradio UI
+# 5.  Launch Gradio Blocks UI
 # ---------------------------------------------------------------------------
 with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue")) as demo:
     gr.Markdown("### SchoolSpirit AI Chat")
@@ -170,19 +175,24 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue")) as demo:
     )
 
     state = gr.State(
-        {"history_raw": [{"role": "system", "content": SYSTEM_MSG}]}
+        {
+            "raw": [
+                {"role": "system", "content": SYSTEM_MSG},
+                {"role": "assistant", "content": WELCOME_MSG},
+            ]
+        }
     )
 
     with gr.Row():
         txt = gr.Textbox(
             placeholder="Type your question here…",
             show_label=False,
-            lines=1,
             scale=4,
+            lines=1,
         )
-        send = gr.Button("Send", variant="primary")
+        send_btn = gr.Button("Send", variant="primary")
 
-    send.click(chat_fn, inputs=[txt, chatbot, state], outputs=[chatbot, state])
+    send_btn.click(chat_fn, inputs=[txt, chatbot, state], outputs=[chatbot, state])
     txt.submit(chat_fn, inputs=[txt, chatbot, state], outputs=[chatbot, state])
 
 demo.launch()