Update app.py

app.py

@@ -1,14 +1,15 @@
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+from threading import Thread
 import torch
 import time
 import psutil
 
-# 1.5B is the "sweet spot" for speed vs intelligence on CPU
+# 1.5B is perfect for CPU stability
 model_id = "unsloth/DeepSeek-R1-Distill-Qwen-1.5B-GGUF"
 filename = "DeepSeek-R1-Distill-Qwen-1.5B-Q4_K_M.gguf"
 
-print("Initializing model...")
+print("Loading stable model...")
 model = AutoModelForCausalLM.from_pretrained(
     model_id,
     gguf_file=filename,
@@ -17,74 +18,67 @@ model = AutoModelForCausalLM.from_pretrained(
 )
 tokenizer = AutoTokenizer.from_pretrained(model_id)
 
-SYSTEM_PROMPT = "You are a math assistant. Think in <think> tags, then output JSON."
-
 def get_stats():
     vm = psutil.virtual_memory()
-    return f"RAM Usage: {vm.percent}% ({vm.used / 1024**3:.1f}GB / {vm.total / 1024**3:.1f}GB)"
+    return f"RAM: {vm.percent}% | {vm.used / 1024**3:.1f}GB / 16GB"
 
 def chat(message, history):
-    # Constructing the prompt
-    prompt = f"system\n{SYSTEM_PROMPT}\nuser\n{message}\nassistant\n<think>\n"
+    # Proper DeepSeek-R1/Qwen prompt format
+    prompt = f"<｜begin_of_sentence｜><｜User｜>{message}<｜Assistant｜><think>\n"
+    
     inputs = tokenizer(prompt, return_tensors="pt").to("cpu")
-    input_ids = inputs.input_ids
+    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
     
+    # Generation kwargs
+    generation_kwargs = dict(
+        inputs,
+        streamer=streamer,
+        max_new_tokens=1024,
+        do_sample=False, # Keeps it fast and logical for math
+        pad_token_id=tokenizer.eos_token_id
+    )
+
+    # Start generation in a background thread
+    thread = Thread(target=model.generate, kwargs=generation_kwargs)
+    thread.start()
+
     start_time = time.time()
-    generated_tokens = 0
-    full_response = ""
+    generated_text = ""
+    token_count = 0
     
-    # Simple streaming generation logic
-    # We use a loop to yield updates to the UI
-    for _ in range(512): # max_new_tokens
-        output = model.generate(
-            input_ids, 
-            max_new_tokens=1, 
-            do_sample=False, 
-            pad_token_id=tokenizer.eos_token_id
-        )
-        
-        new_token_id = output[0][-1]
-        if new_token_id == tokenizer.eos_token_id:
-            break
-            
-        new_token = tokenizer.decode(new_token_id)
-        full_response += new_token
-        generated_tokens += 1
-        
-        # Calculate stats
-        elapsed_time = time.time() - start_time
-        tps = generated_tokens / elapsed_time if elapsed_time > 0 else 0
-        stats = f"⏱️ {elapsed_time:.1f}s | ⚡ {tps:.2f} tokens/s | {get_stats()}"
-        
-        # Prepare for next iteration
-        input_ids = torch.cat([input_ids, output[:, -1:]], dim=-1)
-        
-        yield full_response, stats
+    # Yield from the streamer for real-time UI updates
+    for new_text in streamer:
+        generated_text += new_text
+        token_count += 1
+        elapsed = time.time() - start_time
+        tps = token_count / elapsed if elapsed > 0 else 0
+        stats = f"⏱️ {elapsed:.1f}s | ⚡ {tps:.2f} t/s | {get_stats()}"
+        yield generated_text, stats
 
-with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# 🧠 DeepSeek-R1 CPU Dashboard")
+# Custom Gradio Blocks for better UI
+with gr.Blocks(theme=gr.themes.Default()) as demo:
+    gr.Markdown("# 🚀 DeepSeek-R1 CPU Optimizer")
     
     with gr.Row():
         with gr.Column(scale=4):
-            chatbot = gr.Chatbot(label="Reasoning & JSON Output")
-            msg = gr.Textbox(label="Ask a Math Question", placeholder="e.g., What is the square root of 144 plus 5?")
+            chatbot = gr.Chatbot(label="Response (Thinking + JSON)")
+            msg = gr.Textbox(label="Input", placeholder="Enter your math problem...")
         with gr.Column(scale=1):
-            stats_output = gr.Markdown("### System Stats\nWaiting for input...")
-            clear = gr.Button("Clear Chat")
+            stats_box = gr.Markdown("### Live Stats\nWaiting...")
+            clear = gr.Button("Clear")
 
-    def user_input(user_message, history):
-        return "", history + [[user_message, None]]
+    def respond(message, chat_history):
+        chat_history.append((message, ""))
+        return "", chat_history
 
-    def bot_response(history):
-        user_message = history[-1][0]
-        history[-1][1] = ""
-        # stream the response
-        for chunk, stats in chat(user_message, history[:-1]):
-            history[-1][1] = chunk
-            yield history, stats
+    def stream_bot(chat_history):
+        user_input = chat_history[-1][0]
+        for content, stats in chat(user_input, chat_history[:-1]):
+            chat_history[-1][1] = content
+            yield chat_history, stats
 
-    msg.submit(user_input, [msg, chatbot], [msg, chatbot], queue=False).then(
-        bot_response, chatbot, [chatbot, stats_output]
+    msg.submit(respond, [msg, chatbot], [msg, chatbot]).then(
+        stream_bot, chatbot, [chatbot, stats_box]
     )
     clear.click(lambda: None, None, chatbot, queue=False)