Update app.py

app.py

@@ -1,5 +1,4 @@
 import os
-# SPEED FIX 1: Maximize CPU core usage for Hugging Face Free Tier (2 vCPUs)
 os.environ["OMP_NUM_THREADS"] = "2"
 
 import gradio as gr
@@ -7,76 +6,90 @@ import torch
 import gc
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 from huggingface_hub import hf_hub_download
-from threading import Thread
+from threading import Thread, Event
 
-# SPEED FIX 2: Explicitly tell PyTorch to use both CPU cores
 torch.set_num_threads(2)
 
-model_path = "ruhzi/Indian_History_SLM" 
-
+model_path = "ruhzi/Indian_History_SLM"
 tokenizer = AutoTokenizer.from_pretrained(model_path)
-
 template_file = hf_hub_download(repo_id=model_path, filename="chat_template.jinja")
 with open(template_file, "r", encoding="utf-8") as f:
     tokenizer.chat_template = f.read()
 
-# SPEED FIX 3: Removed device_map and used float32 (Native CPU math is faster)
 model = AutoModelForCausalLM.from_pretrained(
     model_path,
-    torch_dtype=torch.float32, 
+    torch_dtype=torch.float32,
     low_cpu_mem_usage=True
 )
 
+# Global stop event — shared across the current generation
+stop_event = Event()
+
 def chat_inference(message, history):
+    global stop_event
+
+    # Signal any previous generation to stop, then reset for this run
+    stop_event.set()
+    stop_event = Event()
+    current_stop = stop_event  # Capture a local reference for this generation
+
     messages = []
-    
-    # MEMORY PROTECTION: Only keep the last 3 conversational turns
     recent_history = history[-3:] if len(history) > 3 else history
-    
     for user_msg, assistant_msg in recent_history:
         messages.append({"role": "user", "content": user_msg})
         messages.append({"role": "assistant", "content": assistant_msg})
     messages.append({"role": "user", "content": message})
-    
+
     input_text = tokenizer.apply_chat_template(
         messages,
         tokenize=False,
         add_generation_prompt=True,
-        enable_thinking=False 
+        enable_thinking=False
     )
-    
-    # Explicitly send to CPU
+
     inputs = tokenizer([input_text], return_tensors="pt").to("cpu")
 
-    streamer = TextIteratorStreamer(tokenizer, timeout=60.0, skip_prompt=True, skip_special_tokens=True)
+    streamer = TextIteratorStreamer(
+        tokenizer,
+        timeout=60.0,
+        skip_prompt=True,
+        skip_special_tokens=True
+    )
 
     generate_kwargs = dict(
         **inputs,
         streamer=streamer,
-        max_new_tokens=512, # SPEED FIX 4: Kept at 512 for faster, punchier demo responses
+        max_new_tokens=512,
         do_sample=True,
         temperature=0.7,
         top_p=0.8,
     )
 
-    t = Thread(target=model.generate, kwargs=generate_kwargs)
+    t = Thread(target=model.generate, kwargs=generate_kwargs, daemon=True)
     t.start()
 
     partial_message = ""
-    for new_token in streamer:
-        partial_message += new_token
-        yield partial_message
-            
-    # MEMORY PROTECTION: Cleanup after generation finishes
-    del inputs
-    gc.collect()
+    try:
+        for new_token in streamer:
+            if current_stop.is_set():
+                # Drain the streamer so the thread can exit cleanly
+                for _ in streamer:
+                    pass
+                break
+            partial_message += new_token
+            yield partial_message
+    finally:
+        # Always clean up, whether generation finished or was stopped
+        del inputs
+        gc.collect()
+
 
 demo = gr.ChatInterface(
     fn=chat_inference,
     title="Indian History SLM",
     description="Ask me anything about Indian History!",
-    # CRASH PROTECTION: The strict queue. 1 user at a time.
-    concurrency_limit=1 
+    stop_btn="Stop",      # Renders the Stop button
+    concurrency_limit=1,
 )
 
 if __name__ == "__main__":