Update app.py

app.py

@@ -2,54 +2,41 @@ import os, re, time, datetime, traceback, torch, gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 from transformers.utils import logging as hf_logging
 
-# ───────────────── logging ─────────────────────────────────────────
+# ---------------- logging ----------------------------------------------------
 os.environ["HF_HOME"] = "/data/.huggingface"
-LOG = "/data/requests.log"
-
-
 def log(msg: str):
     ts = datetime.datetime.utcnow().strftime("%H:%M:%S.%f")[:-3]
-    line = f"[{ts}] {msg}"
-    print(line, flush=True)
-    try:
-        with open(LOG, "a") as f:
-            f.write(line + "\n")
-    except FileNotFoundError:
-        pass
-
-
-# ───────────────── config ──────────────────────────────────────────
-MODEL_ID = "ibm-granite/granite-3.3-2b-instruct"
-MAX_PAIRS = 4            # user/assistant pairs to keep
-MAX_TOKENS = 128
-MAX_INPUT_CH = 300
-
-SYSTEM_MSG = (
-    "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 supplies GPU servers to K‑12 schools.\n\n"
+    print(f"[{ts}] {msg}", flush=True)
+
+# ---------------- config -----------------------------------------------------
+MODEL_ID            = "ibm-granite/granite-3.3-2b-instruct"
+MAX_TURNS, MAX_TOKS = 4, 64
+SYSTEM_PROMPT = (
+    "You are **SchoolSpirit AI**, digital mascot of SchoolSpirit AI LLC, "
+    "founded by Charles Norton in 2025.  The company installs on‑prem AI chat "
+    "mascots, fine‑tunes LLMs, and ships turnkey GPU hardware to K‑12 schools.\n\n"
     "RULES:\n"
-    "• Friendly, concise (≤ 4 sentences) unless user wants detail.\n"
-    "• If unsure or out of scope, say so and suggest human follow‑up.\n"
-    "• No personal‑data collection, no medical/legal/financial advice.\n"
-    "• Avoid profanity, politics, and mature themes."
+    "• Friendly, concise (≤ 4 sentences) unless asked for detail.\n"
+    "• If unsure or out‑of‑scope: say so & suggest human follow‑up.\n"
+    "• Do NOT collect personal data; no medical/legal/financial advice.\n"
+    "• No profanity, politics, or mature themes."
 )
-WELCOME_MSG = "Welcome to SchoolSpirit AI! Do you have any questions?"
+WELCOME_BUBBLE = "Welcome to SchoolSpirit AI! Do you have any questions?"
 
 strip = lambda s: re.sub(r"\s+", " ", s.strip())
 
 
-# ───────────────── model load (GPU fp16 → CPU) ─────────────────────
+# ---------------- load model (GPU fp16 → CPU) --------------------------------
 hf_logging.set_verbosity_error()
 try:
     tok = AutoTokenizer.from_pretrained(MODEL_ID)
     if torch.cuda.is_available():
-        log("GPU detected → FP16")
+        log("GPU fp16")
         model = AutoModelForCausalLM.from_pretrained(
             MODEL_ID, device_map="auto", torch_dtype=torch.float16
         )
     else:
-        log("CPU fallback")
+        log("CPU fp32")
         model = AutoModelForCausalLM.from_pretrained(
             MODEL_ID, device_map="cpu", torch_dtype="auto", low_cpu_mem_usage=True
         )
@@ -58,93 +45,44 @@ try:
         "text-generation",
         model=model,
         tokenizer=tok,
-        max_new_tokens=MAX_TOKENS,
+        max_new_tokens=MAX_TOKS,
         do_sample=True,
-        temperature=0.65,
+        temperature=0.6,
     )
-    MODEL_ERR = None
-except Exception as exc:  # noqa: BLE001
-    MODEL_ERR, gen = f"Model load error: {exc}", None
-    log(MODEL_ERR)
-
-
-# ───────────────── helper ──────────────────────────────────────────
-def build_prompt(msgs):
-    """Granite likes ### markers"""
-    lines = [f"### System:\n{SYSTEM_MSG}"]
-    for m in msgs:
-        if m["role"] == "user":
-            lines.append(f"### User:\n{m['content']}")
-        elif m["role"] == "assistant":
-            lines.append(f"### Assistant:\n{m['content']}")
-    lines.append("### Assistant:")
-    return "\n".join(lines)
-
+except Exception as e:  # noqa: BLE001
+    gen = None
+    log(f"Model load error: {e}")
 
-def trim(msgs):
-    """Keep system + last MAX_PAIRS*2 messages"""
-    convo = [m for m in msgs if m["role"] != "system"]
-    return [{"role": "system", "content": SYSTEM_MSG}] + convo[-MAX_PAIRS * 2 :]
-
-
-# ───────────────── chat callback ───────────────────────────────────
-def chat_fn(user_msg, history, state):
-    """
-    user_msg : str
-    history  : list[dict] for UI (assistant & user only)
-    state    : {"msgs": full_message_history_with_system}
-    """
-    if MODEL_ERR:
-        history.append({"role": "assistant", "content": MODEL_ERR})
-        return history, state
+# ---------------- chat callback ---------------------------------------------
+def chat_fn(user_msg: str, history: list[tuple[str, str]]):
+    if gen is None:
+        return "Model failed to load. Please try later."
 
     user_msg = strip(user_msg or "")
     if not user_msg:
-        history.append({"role": "assistant", "content": "Please type something."})
-        return history, state
-    if len(user_msg) > MAX_INPUT_CH:
-        history.append(
-            {
-                "role": "assistant",
-                "content": f"Message too long (>{MAX_INPUT_CH} characters).",
-            }
-        )
-        return history, state
+        return "Please type something."
 
-    # Update raw history
-    state["msgs"].append({"role": "user", "content": user_msg})
-    state["msgs"] = trim(state["msgs"])
+    # build prompt: system + last MAX_TURNS pairs + new user msg
+    pairs = history[-MAX_TURNS:]
+    prompt = SYSTEM_PROMPT + "\n"
+    for u, a in pairs:
+        prompt += f"### User:\n{u}\n### Assistant:\n{a}\n"
+    prompt += f"### User:\n{user_msg}\n### Assistant:\n"
 
-    prompt = build_prompt(state["msgs"])
     try:
         raw = gen(prompt)[0]["generated_text"]
         reply = strip(raw.split("### Assistant:", 1)[-1])
+        reply = re.split(r"\b###\s*(?:User|Assistant):", reply, 1)[0].strip()
     except Exception:
-        log("❌ Inference error:\n" + traceback.format_exc())
+        log("generation crash\n" + traceback.format_exc())
         reply = "Sorry—backend crashed. Please try again later."
 
-    # Append to histories
-    state["msgs"].append({"role": "assistant", "content": reply})
-    history.append({"role": "assistant", "content": reply})
-    return history, state
-
-
-# ───────────────── UI ──────────────────────────────────────────────
-with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue")) as demo:
-    chatbot = gr.Chatbot(
-        value=[WELCOME_MSG],
-        label="SchoolSpirit AI",
-        height=480,
-        type="messages",
-    )
-    txt = gr.Textbox(
-        placeholder="Type your question here…",
-        show_label=False,
-        container=False,
-    )
-    state = gr.State({"msgs": [{"role": "system", "content": SYSTEM_MSG}]})
-
-    txt.submit(chat_fn, [txt, chatbot, state], [chatbot, state])
-    txt.submit(lambda _: "", None, txt)  # clear textbox
+    return reply
 
-demo.launch()
+# ---------------- UI ---------------------------------------------------------
+gr.ChatInterface(
+    fn=chat_fn,
+    chatbot=gr.Chatbot(value=[("", WELCOME_BUBBLE)], height=480),
+    title="SchoolSpirit AI Chat",
+    theme=gr.themes.Soft(primary_hue="blue"),
+).launch()