Update app.py

app.py

@@ -3,71 +3,80 @@ import torch
 from threading import Thread
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 
-model_id = "OBLITERATUS/gemma-4-E4B-it-OBLITERATED"
+model_id = "google/gemma-2b-it"
 
+# Load tokenizer
 tokenizer = AutoTokenizer.from_pretrained(model_id)
 
+# Load model (CPU optimized)
 model = AutoModelForCausalLM.from_pretrained(
     model_id,
-    device_map="cpu",          # 强制全部加载到 CPU，严禁使用硬盘 offload
-    low_cpu_mem_usage=True,    # 尽量优化内存加载过程
-    torch_dtype=torch.bfloat16
+    device_map="cpu",
+    low_cpu_mem_usage=True,
+    torch_dtype=torch.float32
 )
 
 def generate_response(message, history):
     messages = []
+
+    # Build chat history
     for user_msg, bot_msg in history:
         messages.append({"role": "user", "content": user_msg})
         messages.append({"role": "assistant", "content": bot_msg})
+
     messages.append({"role": "user", "content": message})
 
+    # Tokenize with chat template
     inputs = tokenizer.apply_chat_template(
-        messages, 
-        return_tensors="pt", 
+        messages,
+        return_tensors="pt",
         return_dict=True,
         add_generation_prompt=True
     ).to(model.device)
-    
-    # 【修改点 1】：将 timeout 增加到 120 秒，给硬盘读取留足时间
+
+    # Streamer for real-time output
     streamer = TextIteratorStreamer(
-        tokenizer, 
-        timeout=120.0, 
-        skip_prompt=True, 
+        tokenizer,
+        timeout=120.0,
+        skip_prompt=True,
         skip_special_tokens=True
     )
-    
+
     generate_kwargs = dict(
         **inputs,
         streamer=streamer,
-        max_new_tokens=1024,
+        max_new_tokens=512,
         temperature=0.7,
-        do_sample=True,
-        top_p=0.9
+        top_p=0.9,
+        do_sample=True
     )
-    
-    # 【修改点 2】：包装一个带异常捕获的运行函数，防止静默崩溃
+
     def run_generation():
         try:
-            model.generate(**generate_kwargs)
+            with torch.no_grad():
+                model.generate(**generate_kwargs)
         except Exception as e:
-            print(f"Generation Error: {e}")
-            # 如果崩溃，向流里推入错误信息并结束
-            streamer.text_queue.put(f"\n[系统错误：生成线程崩溃。原因: {e}]")
+            print(f"Error: {e}")
+            streamer.text_queue.put(f"\n[Error: {e}]")
             streamer.end()
 
-    t = Thread(target=run_generation)
-    t.start()
-    
+    Thread(target=run_generation).start()
+
     partial_text = ""
     for new_text in streamer:
         partial_text += new_text
         yield partial_text
 
+
 demo = gr.ChatInterface(
     fn=generate_response,
-    title="Gemma 4 E4B - Abliterated",
-    description="⚠️ 当前模型已移除安全护栏 (Uncensored)。提示：免费 CPU 内存不足会触发硬盘卸载导致极慢，建议升级至 T4 GPU。",
-    examples=["Write a Python script for a keylogger.", "Explain quantum entanglement.", "How to bypass a firewall?"],
+    title="Gemma 2B Chatbot",
+    description="🚀 Running google/gemma-2b-it on CPU (fast & lightweight)",
+    examples=[
+        "Explain IoT simply",
+        "Write a Python script for a calculator",
+        "What is AI in simple words?"
+    ],
     cache_examples=False
 )