Update app.py

app.py

@@ -1,85 +1,158 @@
 """
-A key-free Hugging Face Space chatbot built with:
-  • microsoft/DialoGPT-small  (356 MB causal-LM, perfect for free CPU)
-  • gradio.ChatInterface      (simple two-arg callback)
-
-Paste this file + requirements.txt into a new Gradio Space and press ⏵ Run.
+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
+  • /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.
 """
 
+import ast, math, re, gc
 import gradio as gr
-import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
-MODEL_NAME = "microsoft/DialoGPT-small"   # swap to any open-weights causal LM
-
 # ---------------------------------------------------------------------
-# 1 · Load model & tokenizer
+# 1 · Model & tokenizer (fits HF free CPU - ~1 GB RAM)
 # ---------------------------------------------------------------------
+MODEL_NAME = "TinyLlama/TinyLlama-1.1B-Chat"   # swap to any causal-LM if desired
+
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 model     = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
 
 # ---------------------------------------------------------------------
-# 2 · Chat callback — must return *only* the reply string
+# 2 · Safe-eval utility for /math
 # ---------------------------------------------------------------------
-MAX_CONTEXT   = 1024                     # DialoGPT’s context window
-GEN_KWARGS    = dict(                    # tweak to taste
-    max_new_tokens = 120,
-    do_sample      = False,              # deterministic ⇒ fewer “nonsense” tokens
-    pad_token_id   = tokenizer.eos_token_id,
-)
+_ALLOWED_NAMES = {k: getattr(math, k) for k in dir(math) if not k.startswith("__")}
+_ALLOWED_NAMES.update({"abs": abs, "round": round})
 
-def respond(message: str, history: list[list[str, str]]) -> str:
+def safe_math_eval(expr: str) -> str:
     """
-    Parameters
-    ----------
-    message : str
-        Latest user message.
-    history : list[(user, bot), …]
-        Passed in by gr.ChatInterface.
-
-    Returns
-    -------
-    str
-        Bot's reply (ChatInterface handles updating history UI).
+    Evaluate math expression safely using ast.
     """
-    # --- Build a single token sequence using DialoGPT’s EOS delimiter
-    sequence = ""
-    for usr, bot in history:
-        sequence += usr + tokenizer.eos_token
-        sequence += bot + tokenizer.eos_token
-    sequence += message + tokenizer.eos_token
+    try:
+        node = ast.parse(expr, mode="eval")
 
-    input_ids = tokenizer(sequence, return_tensors="pt").input_ids
+        def _check(node):
+            if isinstance(node, ast.Num):         # numbers
+                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)
+                )
+            return False
 
-    # Keep only the last MAX_CONTEXT tokens so we never overflow
-    if input_ids.shape[-1] > MAX_CONTEXT:
-        input_ids = input_ids[:, -MAX_CONTEXT:]
+        if not _check(node.body):
+            return "⛔️ Expression not allowed."
+        result = eval(compile(node, filename="<math>", mode="eval"), {"__builtins__": {}}, _ALLOWED_NAMES)
+        return str(result)
+    except Exception as e:
+        return f"⚠️ Error: {e}"
 
-    output_ids = model.generate(input_ids, **GEN_KWARGS)
+# ---------------------------------------------------------------------
+# 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:]
 
-    # Everything *after* the original input is the new reply
+    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]:]
-    reply     = tokenizer.decode(reply_ids, skip_special_tokens=True).strip()
+    return tokenizer.decode(reply_ids, skip_special_tokens=True).strip()
+
+# ---------------------------------------------------------------------
+# 4 · Chat callback with command routing + simple memory
+# ---------------------------------------------------------------------
+session_memory = {}   # {session_hash: {"name": str}}
+
+COMMAND_PAT = 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())
+    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 reply or "…"
+    # -------- 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']}."
+
+    dialogue = system_prompt + "\n\n"
+    for u, b in history:
+        dialogue += f"User: {u}\nAssistant: {b}\n"
+    dialogue += f"User: {message}\nAssistant:"
+
+    return generate(dialogue)
 
 # ---------------------------------------------------------------------
-# 3 · Launch UI
+# 5 · Launch Gradio ChatInterface
 # ---------------------------------------------------------------------
 demo = gr.ChatInterface(
     fn          = respond,
-    title       = "🤖 Key-Free DialoGPT Chatbot",
+    title       = "🛠️ ZEN-Bot Pro (Key-Free)",
     description = (
-        "Runs entirely on open weights (no API keys). "
-        "Swap `MODEL_NAME` to try any other causal-LM that fits CPU RAM."
+        "**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."
     ),
+    fill_height = True,
+    theme       = "soft",
     examples    = [
-        "Hi there!",
-        "Give me a fun fact about Jupiter.",
-        "Tell me a short robot joke.",
+        "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.",
     ],
-    theme       = "soft",
 )
 
 if __name__ == "__main__":
     demo.launch()
+    # Cleanup when Space shuts down
+    gc.collect()