Update app.py

app.py

@@ -1,9 +1,4 @@
-import os
-import re
-import time
-import datetime
-import traceback
-import torch
+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
@@ -13,70 +8,57 @@ from transformers.utils import logging as hf_logging
 # ---------------------------------------------------------------------------
 os.environ["HF_HOME"] = "/data/.huggingface"
 LOG_FILE = "/data/requests.log"
-
-
-def log(msg: str):
+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_FILE, "a") as f:
-            f.write(line + "\n")
-    except FileNotFoundError:
-        pass
-
+    try: open(LOG_FILE,"a").write(line+"\n")
+    except FileNotFoundError: pass
 
 # ---------------------------------------------------------------------------
-# 1.  Configuration constants
+# 1.  Config
 # ---------------------------------------------------------------------------
-MODEL_ID = "ibm-granite/granite-3.3-2b-instruct"
-CONTEXT_TOKENS = 1800
-MAX_NEW_TOKENS = 64
-TEMPERATURE = 0.6
-MAX_INPUT_CH = 300
+MODEL_ID          = "ibm-granite/granite-3.3-2b-instruct"
+CONTEXT_TOKENS    = 1800
+MAX_NEW_TOKENS    = 96
+TEMPERATURE       = 0.5
+MAX_INPUT_CH      = 300
+RATE_LIMIT_N      = 6      # ↲  max messages
+RATE_LIMIT_WINDOW = 60     # ↲  per seconds
 
 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"
+    "SchoolSpirit AI LLC.  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."
+    "• If uncertain, admit it & suggest contacting a human.\n"
+    "• If you can’t answer, politely direct the user to admin@schoolspiritai.com.\n"
+    "• Avoid profanity, politics, or mature themes."
 )
-WELCOME_MSG = "Welcome to SchoolSpirit AI! Do you have any questions?"
+WELCOME_MSG = "Welcome to SchoolSpirit AI! Ask me about our mascots, fine‑tuning, or GPU servers."
 
 strip = lambda s: re.sub(r"\s+", " ", s.strip())
 
-
 # ---------------------------------------------------------------------------
-# 2.  Load tokenizer + model  (GPU FP‑16 → CPU)
+# 2.  Load model
 # ---------------------------------------------------------------------------
 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")
-        model = AutoModelForCausalLM.from_pretrained(
-            MODEL_ID, device_map="auto", torch_dtype=torch.float16
-        )
-    else:
-        log("CPU fallback")
-        model = AutoModelForCausalLM.from_pretrained(
-            MODEL_ID,
-            device_map="cpu",
-            torch_dtype="auto",
-            low_cpu_mem_usage=True,
-        )
-
+    log("Loading tokenizer / model …")
+    tok = AutoTokenizer.from_pretrained(MODEL_ID)
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_ID,
+        device_map="auto" if torch.cuda.is_available() else "cpu",
+        torch_dtype=torch.float16 if torch.cuda.is_available() else "auto",
+        low_cpu_mem_usage=not torch.cuda.is_available(),
+    )
     generator = pipeline(
         "text-generation",
         model=model,
-        tokenizer=tokenizer,
+        tokenizer=tok,
         max_new_tokens=MAX_NEW_TOKENS,
         do_sample=True,
         temperature=TEMPERATURE,
@@ -85,115 +67,88 @@ try:
     MODEL_ERR = None
     log("Model loaded ✔")
 except Exception as exc:
-    MODEL_ERR = f"Model load error: {exc}"
-    generator = None
+    MODEL_ERR, generator = f"Model load error: {exc}", None
     log(MODEL_ERR)
 
+# ---------------------------------------------------------------------------
+# 3.  Rate‑limiter (IP → timestamps list)
+# ---------------------------------------------------------------------------
+VISITS: dict[str,list[float]] = {}
+def allow(ip:str)->bool:
+    now = time.time()
+    times = VISITS.get(ip,[])
+    times = [t for t in times if now - t < RATE_LIMIT_WINDOW]
+    if len(times) >= RATE_LIMIT_N:
+        VISITS[ip] = times  # cleanup stale entries
+        return False
+    times.append(now)
+    VISITS[ip] = times
+    return True
 
 # ---------------------------------------------------------------------------
-# 3.  Helper: build prompt under token budget
+# 4.  Build prompt within token budget
 # ---------------------------------------------------------------------------
-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 build_prompt(raw_history:list[dict])->str:
     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
+        prefix = {"user":"User:","assistant":"AI:"}.get(msg["role"],"")
+        return msg["content"] if not prefix else f"{prefix} {msg['content']}"
+    system = raw_history[0]          # first is system
+    convo  = raw_history[1:]
     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:
+        parts = [system["content"]] + [render(m) for m in convo] + ["AI:"]
+        if len(tok.encode("\n".join(parts), add_special_tokens=False)) <= CONTEXT_TOKENS or len(convo)<=2:
             break
         convo = convo[2:]
-
-    return "\n".join(prompt_parts)
-
+    return "\n".join(parts)
 
 # ---------------------------------------------------------------------------
-# 4.  Chat callback
+# 5.  Chat callback
 # ---------------------------------------------------------------------------
-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
-    """
+def chat_fn(user_msg:str, display_history:list, state:dict, request:gr.Request):
+    ip = request.client.host if request else "unknown"
+    if not allow(ip):
+        reply = "Rate limit exceeded — please wait a minute and try again."
+        display_history.append((user_msg, reply))
+        return display_history, state
+
     user_msg = strip(user_msg or "")
     if not user_msg:
         return display_history, state
-
     if len(user_msg) > MAX_INPUT_CH:
         display_history.append((user_msg, f"Input >{MAX_INPUT_CH} chars."))
         return display_history, state
-
     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})
-
-    # --- Build prompt within token budget
+    state["raw"].append({"role":"user","content":user_msg})
     prompt = build_prompt(state["raw"])
 
-    # --- Generate
     try:
         start = time.time()
-        result = generator(prompt)[0]
-        reply = strip(result["generated_text"])
-        if "User:" in reply:
-            reply = reply.split("User:", 1)[0].strip()
-        log(f"Reply in {time.time() - start:.2f}s ({len(reply)} chars)")
+        reply = strip(generator(prompt)[0]["generated_text"])
+        if "User:" in reply: reply = reply.split("User:",1)[0].strip()
+        log(f"{ip} ok {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 = "Apologies—internal error. Please try again."
 
-    # --- Append assistant reply to both histories
     display_history.append((user_msg, reply))
-    state["raw"].append({"role": "assistant", "content": reply})
+    state["raw"].append({"role":"assistant","content":reply})
     return display_history, state
 
-
 # ---------------------------------------------------------------------------
-# 5.  Launch Gradio Blocks UI
+# 6.  Launch UI
 # ---------------------------------------------------------------------------
 with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue")) as demo:
     gr.Markdown("### SchoolSpirit AI Chat")
-
-    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},
-            ]
-        }
-    )
-
+    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")
-
-    send_btn.click(chat_fn, inputs=[txt, chatbot, state], outputs=[chatbot, state])
-    txt.submit(chat_fn, inputs=[txt, chatbot, state], outputs=[chatbot, state])
+        txt = gr.Textbox(placeholder="Type your question here…", show_label=False, lines=1, scale=4)
+        btn = gr.Button("Send", variant="primary")
+    btn.click(chat_fn, inputs=[txt,chatbot,state], outputs=[chatbot,state])
+    txt.submit(chat_fn, inputs=[txt,chatbot,state], outputs=[chatbot,state])
 
 demo.launch()