Update app.py

app.py

@@ -19,6 +19,10 @@ import gradio as gr
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from huggingface_hub import snapshot_download
+import re
+
+# 🔹 Set torch num threads to max
+torch.set_num_threads(torch.get_num_threads())
 
 # 🔹 Download & load the model from Hugging Face
 model_name = "HyperX-Sen/Qwen-2.5-7B-Reasoning"
@@ -39,36 +43,44 @@ Respond in the following format:
 </answer>
 """
 
-# 🔹 Function to generate response
+# 🔹 Function to extract reasoning and answer
+def extract_response(full_response):
+    reasoning_match = re.search(r"<reasoning>(.*?)</reasoning>", full_response, re.DOTALL)
+    answer_match = re.search(r"<answer>(.*?)</answer>", full_response, re.DOTALL)
+    reasoning = reasoning_match.group(1).strip() if reasoning_match else ""
+    answer = answer_match.group(1).strip() if answer_match else ""
+    return f"<reasoning>\n{reasoning}\n</reasoning>\n<answer>\n{answer}\n</answer>"
+
+# 🔹 Streaming response function
 def chat_response(user_input, top_p, top_k, temperature, max_length):
     messages = [
         {"role": "system", "content": f"{SYSTEM_PROMPT}"},
         {"role": "user", "content": user_input}
     ]
-
-    # 🔹 Format & tokenize input
+    
     input_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
     inputs = tokenizer(input_text, return_tensors="pt").to(model.device)
-
-    # 🔹 Generate response
+    
     with torch.no_grad():
-        output = model.generate(
+        stream = model.generate(
             **inputs,
             max_length=max_length,
             do_sample=True,
             top_p=top_p,
             top_k=top_k,
-            temperature=temperature
+            temperature=temperature,
+            streamer=True
         )
-
-    # 🔹 Decode output
-    response = tokenizer.decode(output[0], skip_special_tokens=True)
-    return response
+    
+    full_response = ""
+    for token in stream:
+        full_response += tokenizer.decode(token, skip_special_tokens=True)
+        yield extract_response(full_response)
 
 # 🔹 Gradio UI
 with gr.Blocks() as demo:
-    gr.Markdown("# 🤖 Qwen-2.5-7B-Reasoning Chatbot")
-
+    gr.Markdown("# 🤖 Qwen-2.5-7B-Reasoning Chatbot (Streaming)")
+    
     with gr.Row():
         chatbot = gr.Textbox(label="Model Response", lines=8, interactive=False)
     
@@ -80,11 +92,11 @@ with gr.Blocks() as demo:
         top_k = gr.Slider(1, 100, value=50, label="Top-k")
         temperature = gr.Slider(0.1, 1.5, value=0.7, label="Temperature")
         max_length = gr.Slider(128, 1024, value=512, label="Max Length")
-
+    
     with gr.Row():
         submit_button = gr.Button("Generate Response")
-
-    submit_button.click(chat_response, inputs=[user_input, top_p, top_k, temperature, max_length], outputs=[chatbot])
+    
+    submit_button.click(chat_response, inputs=[user_input, top_p, top_k, temperature, max_length], outputs=[chatbot], stream=True)
 
 # 🔹 Launch the Gradio app
 demo.launch()