Update app.py

app.py

@@ -14,7 +14,6 @@ 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}"
@@ -22,18 +21,18 @@ def log(msg: str):
     try:
         with open(LOG_FILE, "a") as f:
             f.write(line + "\n")
-    except FileNotFoundError:
+    except Exception:
+        # If logging to file fails, we still want the process to continue
         pass
 
-
 # ---------------------------------------------------------------------------
 # 1.  Configuration constants
 # ---------------------------------------------------------------------------
-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"
+CONTEXT_TOKENS = 1800
 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 "
@@ -50,7 +49,6 @@ 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)
 # ---------------------------------------------------------------------------
@@ -67,10 +65,7 @@ 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(
@@ -80,7 +75,8 @@ try:
         max_new_tokens=MAX_NEW_TOKENS,
         do_sample=True,
         temperature=TEMPERATURE,
-        return_full_text=False,    # ← only return the newly generated text
+        return_full_text=False,
+        streaming=True,    # ← enable token-by-token streaming
     )
     MODEL_ERR = None
     log("Model loaded ✔")
@@ -89,81 +85,96 @@ except Exception as exc:
     generator = None
     log(MODEL_ERR)
 
-
 # ---------------------------------------------------------------------------
-# 3.  Helper: build prompt under token budget
+# 3.  Helper: build prompt under token budget (with fallback on error)
 # ---------------------------------------------------------------------------
 def build_prompt(raw_history: list[dict]) -> str:
-    """
-    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 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"]
-
-    # 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:
-            break
-        # drop oldest user+assistant pair
-        convo = convo[2:]
-
-    return "\n".join(prompt_parts)
+    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)
 
+    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
+# 4.  Chat callback (streaming generator + robust error handling)
 # ---------------------------------------------------------------------------
 def chat_fn(user_msg: str, display_history: list, state: dict):
-    """
-    display_history : list[tuple[str,str]] for UI
-    state["raw"]    : list[dict]          for prompting
-    """
     user_msg = strip(user_msg or "")
+    # Yield nothing if empty input
     if not user_msg:
-        return display_history, state
+        yield display_history, state
+        return
 
+    # Input length check
     if len(user_msg) > MAX_INPUT_CH:
         display_history.append((user_msg, f"Input >{MAX_INPUT_CH} chars."))
-        return display_history, state
+        yield display_history, state
+        return
 
+    # Model load error
     if MODEL_ERR:
         display_history.append((user_msg, MODEL_ERR))
-        return display_history, state
-
-    # --- Update raw history
-    state["raw"].append({"role": "user", "content": user_msg})
+        yield display_history, state
+        return
 
-    # --- Build prompt within token budget
-    prompt = build_prompt(state["raw"])
-
-    # --- Generate
     try:
-        start = time.time()
-        result = generator(prompt)[0]
-        reply = strip(result["generated_text"])
-        # ── NEW: truncate at any hallucinated next "User:" turn
-        if "User:" in reply:
-            reply = reply.split("User:", 1)[0].strip()
-        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."
+        # Add user to history and display
+        state["raw"].append({"role": "user", "content": user_msg})
+        display_history.append((user_msg, ""))
 
-    # --- Append assistant reply to both histories
-    display_history.append((user_msg, reply))
-    state["raw"].append({"role": "assistant", "content": reply})
-    return display_history, state
+        prompt = build_prompt(state["raw"])
+        start = time.time()
+        partial = ""
+
+        # Stream tokens as they arrive
+        for chunk in generator(prompt):
+            try:
+                new_text = strip(chunk.get("generated_text", ""))
+                # Truncate any hallucinated next-turn
+                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:
+                # Skip malformed chunk but keep streaming
+                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)")
 
+    except Exception:
+        # Catch-all for any unexpected errors in chat flow
+        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
 
 # ---------------------------------------------------------------------------
 # 5.  Launch Gradio Blocks UI
@@ -178,24 +189,22 @@ 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},
-            ]
-        }
+        {"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,
-        )
+        txt = gr.Textbox(placeholder="Type your question here…", show_label=False, scale=4, lines=1)
         send_btn = gr.Button("Send", variant="primary")
 
-    send_btn.click(chat_fn, inputs=[txt, chatbot, state], outputs=[chatbot, state])
-    txt.submit(chat_fn, inputs=[txt, chatbot, state], outputs=[chatbot, state])
+    # Enable streaming updates in the UI
+    send_btn.click(chat_fn, inputs=[txt, chatbot, state], outputs=[chatbot, state], stream=True)
+    txt.submit(chat_fn, inputs=[txt, chatbot, state], outputs=[chatbot, state], stream=True)
 
-demo.launch()
+if __name__ == "__main__":
+    try:
+        demo.launch()
+    except Exception as exc:
+        log(f"UI launch error:\n{traceback.format_exc()}")