Update app.py

app.py

@@ -1,15 +1,14 @@
-import os
-import re
-import time
-import datetime
-import traceback
-import torch
+##############################################################################
+# SchoolSpirit AI Chat – full‑context, duplicate‑free implementation
+##############################################################################
+
+import os, re, time, datetime, traceback, torch
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 from transformers.utils import logging as hf_logging
 
 # ---------------------------------------------------------------------------
-# 0.  Paths & basic logging helper
+# 0. Basic logging
 # ---------------------------------------------------------------------------
 os.environ["HF_HOME"] = "/data/.huggingface"
 LOG_FILE = "/data/requests.log"
@@ -27,24 +26,24 @@ def log(msg: str):
 
 
 # ---------------------------------------------------------------------------
-# 1.  Configuration constants
+# 1. Configuration
 # ---------------------------------------------------------------------------
-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
+MODEL_ID = "ibm-granite/granite-3.3-2b-instruct"  # 2‑B local model
+CONTEXT_TOKENS = 1800      # keep prompt below this many tokens
 MAX_NEW_TOKENS = 64
 TEMPERATURE = 0.6
-MAX_INPUT_CH = 300      # UI safeguard
+MAX_INPUT_CH = 300
 
 SYSTEM_MSG = (
-    "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"
+    "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"
     "RULES:\n"
-    "• 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."
+    "• 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."
 )
 WELCOME_MSG = "Welcome to SchoolSpirit AI! Do you have any questions?"
 
@@ -52,20 +51,19 @@ strip = lambda s: re.sub(r"\s+", " ", s.strip())
 
 
 # ---------------------------------------------------------------------------
-# 2.  Load tokenizer + model  (GPU FP‑16 → CPU)
+# 2. Load tokenizer + model (GPU FP‑16 → CPU fallback)
 # ---------------------------------------------------------------------------
 hf_logging.set_verbosity_error()
 try:
-    log("Loading tokenizer …")
     tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
 
     if torch.cuda.is_available():
-        log("GPU detected → FP‑16")
+        log("GPU found → loading model in FP‑16")
         model = AutoModelForCausalLM.from_pretrained(
             MODEL_ID, device_map="auto", torch_dtype=torch.float16
         )
     else:
-        log("CPU fallback")
+        log("No GPU → CPU load")
         model = AutoModelForCausalLM.from_pretrained(
             MODEL_ID,
             device_map="cpu",
@@ -78,12 +76,11 @@ try:
         model=model,
         tokenizer=tokenizer,
         max_new_tokens=MAX_NEW_TOKENS,
-        do_sample=True,
         temperature=TEMPERATURE,
-        return_full_text=False,    # ← only return the newly generated text
+        do_sample=True,
     )
     MODEL_ERR = None
-    log("Model loaded ✔")
+    log("Model ready ✔")
 except Exception as exc:
     MODEL_ERR = f"Model load error: {exc}"
     generator = None
@@ -91,79 +88,77 @@ except Exception as exc:
 
 
 # ---------------------------------------------------------------------------
-# 3.  Helper: build prompt under token budget
+# 3. Prompt builder that respects a token budget
 # ---------------------------------------------------------------------------
-def build_prompt(raw_history: list[dict]) -> str:
+def build_prompt(history_raw: list[dict]) -> str:
     """
-    raw_history: list  [{'role':'system'|'user'|'assistant', 'content': str}, ...]
-    Keeps trimming oldest user/assistant pair until total tokens < CONTEXT_TOKENS
+    Accepts full message list. Drops oldest user+assistant pairs until the
+    prompt’s token length is ≤ CONTEXT_TOKENS.
     """
     def render(msg):
         if msg["role"] == "system":
-            return msg["content"]
+            return f"System: {msg['content']}"
         prefix = "User:" if msg["role"] == "user" else "AI:"
         return f"{prefix} {msg['content']}"
 
-    # 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"]
+    # 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"]
 
-    # iterative trim
     while True:
-        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:
+        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:
             break
-        # drop oldest user+assistant pair
+        # remove oldest user+assistant pair
         convo = convo[2:]
 
     return "\n".join(prompt_parts)
 
 
 # ---------------------------------------------------------------------------
-# 4.  Chat callback
+# 4. Chat callback
 # ---------------------------------------------------------------------------
-def chat_fn(user_msg: str, display_history: list, state: dict):
+def chat_fn(user_text, chat_ui, state):
     """
-    display_history : list[tuple[str,str]] for UI
-    state["raw"]    : list[dict]          for prompting
+    chat_ui : list[(user, assistant)] displayed by gr.Chatbot
+    state   : dict{ 'history_raw': [...] } used for prompt construction
     """
-    user_msg = strip(user_msg or "")
-    if not user_msg:
-        return display_history, state
+    user_text = strip(user_text or "")
+    if not user_text:
+        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 len(user_text) > MAX_INPUT_CH:
+        chat_ui.append((user_text, f"Input exceeds {MAX_INPUT_CH} characters."))
+        return chat_ui, state
 
     if MODEL_ERR:
-        display_history.append((user_msg, MODEL_ERR))
-        return display_history, state
+        chat_ui.append((user_text, MODEL_ERR))
+        return chat_ui, state
 
-    # --- Update raw history
-    state["raw"].append({"role": "user", "content": user_msg})
+    # update raw history (system already present)
+    state["history_raw"].append({"role": "user", "content": user_text})
 
-    # --- Build prompt within token budget
-    prompt = build_prompt(state["raw"])
+    prompt = build_prompt(state["history_raw"])
 
-    # --- Generate
     try:
         start = time.time()
-        result = generator(prompt)[0]
-        reply = strip(result["generated_text"])
-        log(f"Reply in {time.time() - start:.2f}s ({len(reply)} chars)")
+        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)")
     except Exception:
-        log("❌ Inference error:\n" + traceback.format_exc())
-        reply = "Apologies—an internal error occurred. Please try again."
+        log("⚠️  Inference error\n" + traceback.format_exc())
+        reply = "Sorry — an internal error occurred. Please try again."
 
-    # --- Append assistant reply to both histories
-    display_history.append((user_msg, reply))
-    state["raw"].append({"role": "assistant", "content": reply})
-    return display_history, state
+    # append to histories
+    chat_ui.append((user_text, reply))
+    state["history_raw"].append({"role": "assistant", "content": reply})
+    return chat_ui, state
 
 
 # ---------------------------------------------------------------------------
-# 5.  Launch Gradio Blocks UI
+# 5. Gradio UI
 # ---------------------------------------------------------------------------
 with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue")) as demo:
     gr.Markdown("### SchoolSpirit AI Chat")
@@ -175,24 +170,19 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue")) as demo:
     )
 
     state = gr.State(
-        {
-            "raw": [
-                {"role": "system", "content": SYSTEM_MSG},
-                {"role": "assistant", "content": WELCOME_MSG},
-            ]
-        }
+        {"history_raw": [{"role": "system", "content": SYSTEM_MSG}]}
     )
 
     with gr.Row():
         txt = gr.Textbox(
             placeholder="Type your question here…",
             show_label=False,
-            scale=4,
             lines=1,
+            scale=4,
         )
-        send_btn = gr.Button("Send", variant="primary")
+        send = gr.Button("Send", variant="primary")
 
-    send_btn.click(chat_fn, inputs=[txt, chatbot, state], outputs=[chatbot, state])
+    send.click(chat_fn, inputs=[txt, chatbot, state], outputs=[chatbot, state])
     txt.submit(chat_fn, inputs=[txt, chatbot, state], outputs=[chatbot, state])
 
 demo.launch()