Update app.py

app.py

@@ -1,13 +1,9 @@
-##############################################################################
-#  SchoolSpirit AI – Full‑context Chatbot Space
-#  -------------------------------------------------
-#  • Maintains a rolling prompt that only truncates when token budget exceeded
-#  • GPU FP‑16 load, CPU fallback
-#  • Detailed logging, duplicate‑proof welcome
-#  • Gradio Blocks UI (textbox + Send button)
-##############################################################################
-
-import os, re, time, datetime, traceback, torch, json, math
+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
@@ -48,7 +44,7 @@ SYSTEM_MSG = (
     "• 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."
+    "• avoid profanity, politics, mature themes."
 )
 WELCOME_MSG = "Welcome to SchoolSpirit AI! Do you have any questions?"
 
@@ -84,6 +80,7 @@ try:
         max_new_tokens=MAX_NEW_TOKENS,
         do_sample=True,
         temperature=TEMPERATURE,
+        return_full_text=False,    # ← only return the newly generated text
     )
     MODEL_ERR = None
     log("Model loaded ✔")
@@ -152,10 +149,9 @@ def chat_fn(user_msg: str, display_history: list, state: dict):
     # --- Generate
     try:
         start = time.time()
-        completion = generator(prompt)[0]["generated_text"]
-        reply = strip(completion.split("AI:", 1)[-1])
-        reply = re.split(r"\b(?:User:|AI:)", reply, 1)[0].strip()
-        log(f"Reply 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 = "Apologies—an internal error occurred. Please try again."