Update app.py

app.py

@@ -6,146 +6,170 @@ Skills
   • /math <expr>           — safe calculator
   • /summarize <text>      — 2-sentence TL;DR
   • /translate_es <text>   — English → Spanish
-  • /ascii <text>          — big FIGlet ASCII art
-  • Any question ending with “?”  → live Wikipedia answer
-Model: facebook/blenderbot-400M-distill  (public,  ~720 MB weights)
+  • /ascii <text>          — FIGlet ASCII art
+  • Any question ending with “?” → live Wikipedia answer
+     ↳ remembers the last topic, so “Who discovered it?” now works.
+Model: facebook/blenderbot-400M-distill  (~720 MB, public)
+
+Copy-paste this file into your Space, keep the same requirements.txt, commit, done.
 """
 
 import ast, math, re, gc, traceback, torch, wikipedia, pyfiglet
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 
-# ──────────────────────────────────────────────────────────────────────────────
+# ────────────────────────────────────────────────────────────────
 # 0 · helpers
-# ──────────────────────────────────────────────────────────────────────────────
-def err(e: Exception) -> str:
-    print("\n=== ZEN-Bot ERROR ===")
+# ────────────────────────────────────────────────────────────────
+def log_err(e: Exception) -> str:
+    print("=== ZEN-Bot ERROR =====================")
     traceback.print_exc()
-    print("====================\n")
+    print("=======================================\n")
     return f"⚠️ {type(e).__name__}: {e}"
 
-# ──────────────────────────────────────────────────────────────────────────────
-# 1 · model (fits free CPU tier)
-# ──────────────────────────────────────────────────────────────────────────────
-MODEL_NAME = "facebook/blenderbot-400M-distill"      # always public
+# ────────────────────────────────────────────────────────────────
+# 1 · model (public, small enough for free CPU/basic)
+# ────────────────────────────────────────────────────────────────
+MODEL = "facebook/blenderbot-400M-distill"
 
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model     = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)
+tokenizer = AutoTokenizer.from_pretrained(MODEL)
+model     = AutoModelForSeq2SeqLM.from_pretrained(MODEL)
 model.eval();  torch.set_grad_enabled(False)
 
-GEN_KW = dict(
-    max_new_tokens = 64,
-    do_sample      = False,                       # fast & deterministic
-)
+GEN_ARGS = dict(max_new_tokens=64, do_sample=False)   # greedy, fast
 
 def llm(prompt: str) -> str:
+    """Generate with BlenderBot-distill, catching any runtime errors."""
     try:
-        ids_in   = tokenizer(prompt, return_tensors="pt",
-                              truncation=True, max_length=1024).input_ids
+        inputs = tokenizer(
+            prompt, return_tensors="pt",
+            truncation=True, max_length=1024
+        ).input_ids
         with torch.no_grad():
-            ids_out = model.generate(ids_in, **GEN_KW)
-        reply_ids = ids_out[0, ids_in.shape[-1]:]
-        return tokenizer.decode(reply_ids, skip_special_tokens=True).strip()
+            out = model.generate(inputs, **GEN_ARGS)
+        answer_ids = out[0, inputs.shape[-1]:]
+        return tokenizer.decode(answer_ids,
+                                skip_special_tokens=True).strip()
     except Exception as e:
-        return err(e)
+        return log_err(e)
 
-# ──────────────────────────────────────────────────────────────────────────────
-# 2 · /math (safe)
-# ──────────────────────────────────────────────────────────────────────────────
-_ALLOWED = {k: getattr(math, k) for k in dir(math) if not k.startswith("__")}
+# ────────���───────────────────────────────────────────────────────
+# 2 · safe /math
+# ────────────────────────────────────────────────────────────────
+_ALLOWED = {k: getattr(math, k) for k in dir(math) if not k.startswith("_")}
 _ALLOWED.update({"abs": abs, "round": round})
 
 def safe_math(expr: str) -> str:
     try:
         node = ast.parse(expr, mode="eval")
+
         def ok(n):
-            if isinstance(n, ast.Num):            return True
-            if isinstance(n, ast.UnaryOp):        return ok(n.operand)
-            if isinstance(n, ast.BinOp):          return ok(n.left) and ok(n.right)
-            if isinstance(n, ast.Call):
-                return (isinstance(n.func, ast.Name)
-                        and n.func.id in _ALLOWED
-                        and all(ok(a) for a in n.args))
-            return False
+            match n:
+                case ast.Num():                    return True
+                case ast.UnaryOp():                return ok(n.operand)
+                case ast.BinOp():                  return ok(n.left) and ok(n.right)
+                case ast.Call():
+                    return (isinstance(n.func, ast.Name)
+                            and n.func.id in _ALLOWED
+                            and all(ok(a) for a in n.args))
+                case _:                            return False
+
         if not ok(node.body):
-            return "⛔️ Expression not allowed."
+            return "⛔️ Only basic maths & trig/log functions allowed."
         return str(eval(compile(node, "<expr>", "eval"),
                         {"__builtins__": {}}, _ALLOWED))
     except Exception as e:
-        return err(e)
+        return log_err(e)
 
-# ──────────────────────────────────────────────────────────────────────────────
-# 3 · Wikipedia Q&A
-# ──────────────────────────────────────────────────────────────────────────────
-def wiki_answer(q: str) -> str | None:
+# ────────────────────────────────────────────────────────────────
+# 3 · Wikipedia Q&A (with last-topic memory)
+# ────────────────────────────────────────────────────────────────
+last_topic: dict[str, str] = {}   # {session_hash: "Penicillin"}
+
+def wiki_answer(q: str, session_id: str) -> str | None:
+    """Return a 3-sentence Wikipedia summary or None if not found."""
     try:
+        # If the question uses 'it', swap with last remembered topic
+        query = q.strip()
+        if last_topic.get(session_id) and re.search(r"\bit\b", query, re.I):
+            query = re.sub(r"\bit\b", last_topic[session_id], query, flags=re.I)
+
         wikipedia.set_lang("en")
-        page  = wikipedia.page(q, auto_suggest=True, redirect=True)
-        summ  = wikipedia.summary(page.title, sentences=3, auto_suggest=False)
-        return f"**{page.title}** — {summ}"
+        page = wikipedia.page(query, auto_suggest=True, redirect=True)
+        last_topic[session_id] = page.title     # remember for pronouns
+        summary = wikipedia.summary(
+            page.title, sentences=3, auto_suggest=False
+        )
+        return f"**{page.title}** — {summary}"
     except (wikipedia.DisambiguationError, wikipedia.PageError):
-        return None
+        return None           # let LLM try instead
     except Exception as e:
-        return err(e)
+        return log_err(e)
 
-# ──────────────────────────────────────────────────────────────────────────────
+# ────────────────────────────────────────────────────────────────
 # 4 · ASCII art
-# ──────────────────────────────────────────────────────────────────────────────
+# ────────────────────────────────────────────────────────────────
 def ascii_art(txt: str) -> str:
-    try:    return pyfiglet.figlet_format(txt, width=120)
-    except Exception as e:  return err(e)
+    try:
+        return pyfiglet.figlet_format(txt, width=120)
+    except Exception as e:
+        return log_err(e)
 
-# ──────────────────────────────────────────────────────────────────────────────
-# 5 · main router
-# ──────────────────────────────────────────────────────────────────────────────
+# ────────────────────────────────────────────────────────────────
+# 5 · router ↔ ChatInterface callback
+# ────────────────────────────────────────────────────────────────
 CMD = re.compile(r"^/(math|summarize|translate_es|ascii)\s+(.+)", re.S | re.I)
 
-def respond(msg: str, hist: list[list[str, str]]) -> str:
-    # 5-A · commands
-    m = CMD.match(msg.strip())
-    if m:
+def respond(message: str,
+            history: list[list[str, str]],
+            request: gr.Request) -> str:
+    session_id = request.session_hash or "anon"
+
+    # 5-A · Command shortcuts
+    if (m := CMD.match(message.strip())):
         cmd, body = m.group(1).lower(), m.group(2).strip()
         if cmd == "math":         return safe_math(body)
         if cmd == "ascii":        return ascii_art(body)
-        if cmd == "summarize":    return llm(f"Summarize in two concise sentences:\n\n{body}\n\nSummary:")
+        if cmd == "summarize":    return llm(f"Summarise in two concise sentences:\n\n{body}\n\nSummary:")
         if cmd == "translate_es": return llm(f"Translate into Spanish (natural):\n\n{body}\n\nSpanish:")
 
-    # 5-B · Wikipedia for factual questions
-    if msg.strip().endswith("?") and len(msg.split()) > 2:
-        wiki = wiki_answer(msg)
+    # 5-B · Wikipedia factual Q&A
+    if message.strip().endswith("?") and len(message.split()) > 2:
+        wiki = wiki_answer(message, session_id)
         if wiki:  return wiki
 
-    # 5-C · normal chat (keep last 6 turns)
-    prompt = "You are ZEN-Bot, a concise, helpful tutor for young AI pioneers.\n\n"
-    for u, b in hist[-6:]:
+    # 5-C · Normal chat  (keep last 6 turns for speed)
+    prompt = ("You are ZEN-Bot, a concise, friendly tutor for young AI pioneers.\n\n")
+    for u, b in history[-6:]:
         prompt += f"User: {u}\nAssistant: {b}\n"
-    prompt += f"User: {msg}\nAssistant:"
+    prompt += f"User: {message}\nAssistant:"
     return llm(prompt)
 
-# ──────────────────────────────────────────────────────────────────────────────
+# ────────────────────────────────────────────────────────────────
 # 6 · Gradio UI
-# ──────────────────────────────────────────────────────────────────────────────
+# ────────────────────────────────────────────────────────────────
 demo = gr.ChatInterface(
     fn          = respond,
     title       = "🚀 ZEN-Bot Ultimate (Key-Free)",
     description = (
         "**Commands**  \n"
-        "• normal chat  \n"
+        "• Normal chat  \n"
         "• `/math 2**5 / (sin(0.5)+1)`  \n"
         "• `/summarize <text>`  \n"
         "• `/translate_es Hello!`  \n"
         "• `/ascii ZEN`  \n"
-        "• Ask any factual question ending with `?` for live Wikipedia answer"
+        "• Ask factual questions ending with `?` (remembers topic for pronouns)"
     ),
-    theme       = "soft",
-    fill_height = True,
     examples    = [
         "Who discovered penicillin?",
         "/ascii AI ROCKS",
         "/math sqrt(144)+log(100,10)",
         "/summarize The Industrial Revolution began in Britain...",
         "/translate_es Good evening, friends!",
+        "Who discovered it?",
     ],
+    fill_height = True,
+    theme       = "soft",
 )
 
 if __name__ == "__main__":