Update app.py

app.py

@@ -14,6 +14,7 @@ from transformers.utils import logging as hf_logging
 os.environ["HF_HOME"] = "/data/.huggingface"
 LOG_FILE = "/data/requests.log"
 
+
 def log(msg: str):
     ts = datetime.datetime.utcnow().strftime("%H:%M:%S.%f")[:-3]
     line = f"[{ts}] {msg}"
@@ -21,8 +22,9 @@ def log(msg: str):
     try:
         with open(LOG_FILE, "a") as f:
             f.write(line + "\n")
-    except Exception:
-        pass  # tolerate logging failures
+    except FileNotFoundError:
+        pass
+
 
 # ---------------------------------------------------------------------------
 # 1.  Configuration constants
@@ -48,6 +50,7 @@ WELCOME_MSG = "Welcome to SchoolSpirit AI! Do you have any questions?"
 
 strip = lambda s: re.sub(r"\s+", " ", s.strip())
 
+
 # ---------------------------------------------------------------------------
 # 2.  Load tokenizer + model  (GPU FP‑16 → CPU)
 # ---------------------------------------------------------------------------
@@ -64,7 +67,10 @@ try:
     else:
         log("CPU fallback")
         model = AutoModelForCausalLM.from_pretrained(
-            MODEL_ID, device_map="cpu", torch_dtype="auto", low_cpu_mem_usage=True
+            MODEL_ID,
+            device_map="cpu",
+            torch_dtype="auto",
+            low_cpu_mem_usage=True,
         )
 
     generator = pipeline(
@@ -75,7 +81,6 @@ try:
         do_sample=True,
         temperature=TEMPERATURE,
         return_full_text=False,
-        streaming=True,    # ← enable token-by-token streaming
     )
     MODEL_ERR = None
     log("Model loaded ✔")
@@ -84,121 +89,88 @@ except Exception as exc:
     generator = None
     log(MODEL_ERR)
 
+
 # ---------------------------------------------------------------------------
-# 3.  Helper: build prompt under token budget (with fallback on error)
+# 3.  Helper: build prompt under token budget
 # ---------------------------------------------------------------------------
 def build_prompt(raw_history: list[dict]) -> str:
-    try:
-        def render(msg):
-            if msg["role"] == "system":
-                return msg["content"]
-            prefix = "User:" if msg["role"] == "user" else "AI:"
-            return f"{prefix} {msg['content']}"
-
-        system_msg = next(m for m in raw_history if m["role"] == "system")
-        convo = [m for m in raw_history if m["role"] != "system"]
-
-        while True:
-            parts = [system_msg["content"]] + [render(m) for m in convo] + ["AI:"]
-            token_len = len(tokenizer.encode("\n".join(parts), add_special_tokens=False))
-            if token_len <= CONTEXT_TOKENS or len(convo) <= 2:
-                break
-            convo = convo[2:]
-        return "\n".join(parts)
+    system_msg = [m for m in raw_history if m["role"] == "system"][0]
+    convo = [m for m in raw_history if m["role"] != "system"]
+
+    def render(msg):
+        prefix = "User:" if msg["role"] == "user" else "AI:"
+        return f"{prefix} {msg['content']}" if msg["role"] != "system" else msg["content"]
+
+    while True:
+        parts = [system_msg["content"]] + [render(m) for m in convo] + ["AI:"]
+        token_len = len(tokenizer.encode("\n".join(parts), add_special_tokens=False))
+        if token_len <= CONTEXT_TOKENS or len(convo) <= 2:
+            break
+        convo = convo[2:]
+    return "\n".join([system_msg["content"]] + [render(m) for m in convo] + ["AI:"])
 
-    except Exception:
-        log("Error building prompt:\n" + traceback.format_exc())
-        # Fallback: include system + last two messages
-        sys_text = next((m["content"] for m in raw_history if m["role"]=="system"), "")
-        tail = [m for m in raw_history if m["role"]!="system"][-2:]
-        fallback = sys_text + "\n" + "\n".join(
-            f"{'User:' if m['role']=='user' else 'AI:'} {m['content']}" for m in tail
-        ) + "\nAI:"
-        return fallback
 
 # ---------------------------------------------------------------------------
-# 4.  Chat callback (streaming generator + robust error handling)
+# 4.  Chat callback with immediate user echo & spinner
 # ---------------------------------------------------------------------------
 def chat_fn(user_msg: str, display_history: list, state: dict):
     user_msg = strip(user_msg or "")
     if not user_msg:
-        yield display_history, state
-        return
+        return display_history, state
 
     if len(user_msg) > MAX_INPUT_CH:
         display_history.append((user_msg, f"Input >{MAX_INPUT_CH} chars."))
-        yield display_history, state
-        return
+        return display_history, state
 
     if MODEL_ERR:
         display_history.append((user_msg, MODEL_ERR))
-        yield display_history, state
-        return
+        return display_history, state
 
-    try:
-        # record user
-        state["raw"].append({"role": "user", "content": user_msg})
-        display_history.append((user_msg, ""))
+    # Immediately append the user message with a placeholder
+    display_history.append((user_msg, ""))
 
-        prompt = build_prompt(state["raw"])
-        start = time.time()
-        partial = ""
-
-        # stream chunks
-        for chunk in generator(prompt):
-            try:
-                new_text = strip(chunk.get("generated_text", ""))
-                if "User:" in new_text:
-                    new_text = new_text.split("User:", 1)[0].strip()
-                partial += new_text
-                display_history[-1] = (user_msg, partial)
-                yield display_history, state
-            except Exception:
-                log("Malformed chunk:\n" + traceback.format_exc())
-                continue
-
-        # finalize
-        full_reply = display_history[-1][1]
-        state["raw"].append({"role": "assistant", "content": full_reply})
-        log(f"Reply in {time.time() - start:.2f}s ({len(full_reply)} chars)")
+    # Update raw history for prompt
+    state["raw"].append({"role": "user", "content": user_msg})
+    prompt = build_prompt(state["raw"])
 
+    # Generate the bot reply
+    try:
+        start = time.time()
+        out = generator(prompt)[0]["generated_text"].strip()
+        # Truncate any hallucinated next "User:"
+        if "User:" in out:
+            out = out.split("User:", 1)[0].strip()
+        reply = out
+        log(f"Reply in {time.time()-start:.2f}s ({len(reply)} chars)")
     except Exception:
-        log("Unexpected chat_fn error:\n" + traceback.format_exc())
-        err = "Apologies—an internal error occurred. Please try again."
-        display_history[-1] = (user_msg, err)
-        state["raw"].append({"role": "assistant", "content": err})
-        yield display_history, state
+        log("❌ Inference error:\n" + traceback.format_exc())
+        reply = "Apologies—an internal error occurred. Please try again."
+
+    # Replace the placeholder with the actual reply
+    display_history[-1] = (user_msg, reply)
+    state["raw"].append({"role": "assistant", "content": reply})
+
+    return display_history, state
+
 
 # ---------------------------------------------------------------------------
-# 5.  Launch Gradio Blocks UI
+# 5.  Launch Gradio Blocks UI with spinner
 # ---------------------------------------------------------------------------
 with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue")) as demo:
     gr.Markdown("### SchoolSpirit AI Chat")
 
-    chatbot = gr.Chatbot(
-        value=[{"role":"assistant","content":WELCOME_MSG}],
-        height=480,
-        label="SchoolSpirit AI",
-        type="messages",  # use the new messages format
-    )
-
-    state = gr.State(
-        {"raw": [
-            {"role": "system", "content": SYSTEM_MSG},
-            {"role": "assistant", "content": WELCOME_MSG},
-        ]}
-    )
+    chatbot = gr.Chatbot(value=[("", WELCOME_MSG)], height=480, label="SchoolSpirit AI")
+    state = gr.State({"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, scale=4, lines=1)
         send_btn = gr.Button("Send", variant="primary")
 
-    # Use streaming=True (not stream) per Gradio API
-    send_btn.click(chat_fn, inputs=[txt, chatbot, state], outputs=[chatbot, state], streaming=True)
-    txt.submit(chat_fn, inputs=[txt, chatbot, state], outputs=[chatbot, state], streaming=True)
+    # show_progress=True displays a spinner while waiting
+    send_btn.click(chat_fn, inputs=[txt, chatbot, state], outputs=[chatbot, state], show_progress=True)
+    txt.submit(chat_fn, inputs=[txt, chatbot, state], outputs=[chatbot, state], show_progress=True)
 
-if __name__ == "__main__":
-    try:
-        demo.launch()
-    except Exception:
-        log("UI launch error:\n" + traceback.format_exc())
+demo.launch()