Update app.py

app.py

@@ -1,158 +1,133 @@
 """
-Advanced, key-free chatbot for Hugging Face Spaces
--------------------------------------------------
-Features
-  • Natural chat with TinyLlama-1.1B-Chat (open weights, ~1 GB)
-  • /math        – secure calculator (basic math + trig/log)
-  • /summarize   – 2-sentence TL;DR
+ZEN-Bot Lite  –  Key-free chatbot for HF Spaces (free CPU)
+Skills
+  • Natural chat
+  • /math        – safe calculator
+  • /summarize   – 2-sentence abstract
   • /translate_es – English → Spanish
-  • Remembers user's name inside the session
-Everything runs through ONE language model, so it stays within the free CPU tier.
+Only one open-weights model (facebook/blenderbot-400M-distill).
 """
 
 import ast, math, re, gc
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 
 # ---------------------------------------------------------------------
-# 1 · Model & tokenizer (fits HF free CPU - ~1 GB RAM)
+# 1 · Model small enough for free CPU
 # ---------------------------------------------------------------------
-MODEL_NAME = "TinyLlama/TinyLlama-1.1B-Chat"   # swap to any causal-LM if desired
+MODEL_NAME = "facebook/blenderbot-400M-distill"   # ~720 MB on disk, loads in ≈3 GB RAM
 
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model     = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
+model     = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)
 
 # ---------------------------------------------------------------------
-# 2 · Safe-eval utility for /math
+# 2 · Safe /math evaluator  (same as before)
 # ---------------------------------------------------------------------
-_ALLOWED_NAMES = {k: getattr(math, k) for k in dir(math) if not k.startswith("__")}
-_ALLOWED_NAMES.update({"abs": abs, "round": round})
+_ALLOWED = {k: getattr(math, k) for k in dir(math) if not k.startswith("__")}
+_ALLOWED.update({"abs": abs, "round": round})
 
-def safe_math_eval(expr: str) -> str:
-    """
-    Evaluate math expression safely using ast.
-    """
+def safe_math(expr: str) -> str:
     try:
         node = ast.parse(expr, mode="eval")
 
-        def _check(node):
-            if isinstance(node, ast.Num):         # numbers
+        def ok(n):
+            if isinstance(n, ast.Num):
                 return True
-            if isinstance(node, ast.BinOp):       # +, -, *, /, **, etc.
-                return _check(node.left) and _check(node.right)
-            if isinstance(node, ast.UnaryOp):     # -1
-                return _check(node.operand)
-            if isinstance(node, ast.Call):        # sin(0.5)
-                return (
-                    isinstance(node.func, ast.Name)
-                    and node.func.id in _ALLOWED_NAMES
-                    and all(_check(arg) for arg in node.args)
-                )
+            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
 
-        if not _check(node.body):
-            return "⛔️ Expression not allowed."
-        result = eval(compile(node, filename="<math>", mode="eval"), {"__builtins__": {}}, _ALLOWED_NAMES)
-        return str(result)
+        if not ok(node.body):
+            return "⛔️ Expression not allowed"
+        return str(eval(compile(node, "<expr>", "eval"),
+                        {"__builtins__": {}}, _ALLOWED))
     except Exception as e:
         return f"⚠️ Error: {e}"
 
 # ---------------------------------------------------------------------
 # 3 · Generation helper
 # ---------------------------------------------------------------------
-MAX_NEW_TOKENS = 160
-TOKEN_LIMIT    = 1024          # truncate long histories
-
-def generate(prompt: str) -> str:
-    input_ids = tokenizer(prompt, return_tensors="pt").input_ids
-    if input_ids.shape[-1] > TOKEN_LIMIT:
-        input_ids = input_ids[:, -TOKEN_LIMIT:]
-
-    output_ids = model.generate(
-        input_ids,
-        max_new_tokens=MAX_NEW_TOKENS,
-        do_sample=True,
-        top_p=0.92,
-        temperature=0.7,
-        pad_token_id=tokenizer.eos_token_id,
-    )
-    reply_ids = output_ids[0, input_ids.shape[-1]:]
+GEN_KW = dict(
+    max_new_tokens = 128,
+    do_sample      = True,
+    top_p          = 0.92,
+    temperature    = 0.7,
+)
+
+def llm(prompt: str) -> str:
+    input_ids  = tokenizer(prompt, return_tensors="pt").input_ids
+    output_ids = model.generate(**GEN_KW, input_ids=input_ids)
+    reply_ids  = output_ids[0, input_ids.shape[-1]:]
     return tokenizer.decode(reply_ids, skip_special_tokens=True).strip()
 
 # ---------------------------------------------------------------------
-# 4 · Chat callback with command routing + simple memory
+# 4 · Regex router + lightweight name memory (stored inside history)
 # ---------------------------------------------------------------------
-session_memory = {}   # {session_hash: {"name": str}}
-
-COMMAND_PAT = re.compile(r"^/(math|summarize|translate_es)\s+(.*)", re.S | re.I)
+CMD = re.compile(r"^/(math|summarize|translate_es)\s+(.+)", re.S | re.I)
 
-def respond(message: str, history: list[list[str, str]], session: gr.Request) -> str:
-    sess_id = session.session_hash or "anon"
-    mem     = session_memory.setdefault(sess_id, {})
-
-    # -------- handle special commands --------
-    m = COMMAND_PAT.match(message.strip())
+def reply(message: str, history: list[list[str, str]]) -> str:
+    # 4.1 · Command handling first
+    m = CMD.match(message.strip())
     if m:
         cmd, payload = m.group(1).lower(), m.group(2).strip()
         if cmd == "math":
-            return safe_math_eval(payload)
-        elif cmd == "summarize":
-            prompt = (
-                "Summarize the following text in 2 concise sentences:\n\n"
-                f"{payload}\n\nSummary:"
-            )
-            return generate(prompt)
-        elif cmd == "translate_es":
-            prompt = (
-                "Translate the following text from English to Spanish (keep it natural):\n\n"
-                f"{payload}\n\nSpanish:"
-            )
-            return generate(prompt)
-
-    # -------- name capture (very lightweight memory) --------
-    name_match = re.search(r"\bmy name is (\w+)", message, re.I)
-    if name_match:
-        mem["name"] = name_match.group(1).capitalize()
+            return safe_math(payload)
+
+        if cmd == "summarize":
+            prompt = ( "Summarize in two concise sentences:\n\n"
+                       f"{payload}\n\nSummary:" )
+            return llm(prompt)
 
-    # -------- regular chat --------
-    system_prompt = (
-        "You are ZEN-Bot, a kind, concise AI assistant for young tech pioneers."
-    )
-    if "name" in mem:
-        system_prompt += f" The user's name is {mem['name']}."
+        if cmd == "translate_es":
+            prompt = ( "Translate the following text from English to Spanish "
+                       "(natural, not literal):\n\n"
+                       f"{payload}\n\nSpanish:" )
+            return llm(prompt)
 
-    dialogue = system_prompt + "\n\n"
+    # 4.2 · Grab user’s name if volunteered
+    name_match = re.search(r"\bmy name is (\w+)", message, re.I)
+    name = name_match.group(1).capitalize() if name_match else None
+
+    # 4.3 · Build dialogue prompt
+    system = "You are ZEN-Bot, a concise, friendly tutor for young AI pioneers."
+    if name:
+        system += f" The user's name is {name}."
+    prompt = system + "\n\n"
     for u, b in history:
-        dialogue += f"User: {u}\nAssistant: {b}\n"
-    dialogue += f"User: {message}\nAssistant:"
+        prompt += f"User: {u}\nAssistant: {b}\n"
+    prompt += f"User: {message}\nAssistant:"
 
-    return generate(dialogue)
+    return llm(prompt)
 
 # ---------------------------------------------------------------------
-# 5 · Launch Gradio ChatInterface
+# 5 · Gradio UI
 # ---------------------------------------------------------------------
 demo = gr.ChatInterface(
-    fn          = respond,
-    title       = "🛠️ ZEN-Bot Pro (Key-Free)",
+    fn          = reply,
+    title       = "🔧 ZEN-Bot Lite (Key-Free)",
     description = (
-        "**Skills**\n"
-        "• Chat naturally\n"
-        "• `/math 1+2*3` – calculator\n"
-        "• `/summarize <text>` – 2-sentence TL;DR\n"
-        "• `/translate_es <text>` – English→Spanish\n\n"
-        "Runs on open weights (TinyLlama-1.1B-Chat) – no API keys needed."
+        "**Try me**\n"
+        "• normal chat\n"
+        "• `/math 3*(sin(0.5)+1)`\n"
+        "• `/summarize Industrial Revolution text...`\n"
+        "• `/translate_es Good morning!`"
     ),
-    fill_height = True,
-    theme       = "soft",
     examples    = [
-        "Hi, my name is Alex!",
-        "/math sin(0.5) ** 2 + cos(0.5) ** 2",
-        "/summarize The James Webb Space Telescope is the most powerful…",
-        "/translate_es Artificial intelligence will change the world.",
+        "Hi, my name is Jordan!",
+        "/math 2**6 / 3",
+        "/summarize The water cycle begins when...",
+        "/translate_es We are building awesome AI projects.",
     ],
+    theme       = "soft",
+    fill_height = True,
 )
 
 if __name__ == "__main__":
     demo.launch()
-    # Cleanup when Space shuts down
     gc.collect()