Update app.py

app.py

@@ -6,34 +6,22 @@ import spaces
 import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 
-#subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
-
 MODEL_ID = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"
-CHAT_TEMPLATE = "َAuto"
 MODEL_NAME = MODEL_ID.split("/")[-1]
 CONTEXT_LENGTH = 4096
 
-
 def predict(message, history, system_prompt, temperature, max_new_tokens, top_k, repetition_penalty, top_p):
-    # Format history with a given chat template
-    
-    
-    stop_tokens = ["<|endoftext|>", "<|im_end|>","|im_end|"]
+    stop_tokens = ["<|endoftext|>", "<|im_end|>", "|im_end|"]
     instruction = '<|im_start|>system\n' + system_prompt + '\n<|im_end|>\n'
+    
     for user, assistant in history:
         instruction += f'<|im_start|>user\n{user}\n<|im_end|>\n<|im_start|>assistant\n{assistant}\n<|im_end|>\n'
     instruction += f'<|im_start|>user\n{message}\n<|im_end|>\n<|im_start|>assistant\n'
     
-    print(instruction)
-    
     streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
-    enc = tokenizer(instruction, return_tensors="pt", padding=True, truncation=True)
+    enc = tokenizer(instruction, return_tensors="pt", truncation=True, max_length=CONTEXT_LENGTH)
     input_ids, attention_mask = enc.input_ids, enc.attention_mask
 
-    if input_ids.shape[1] > CONTEXT_LENGTH:
-        input_ids = input_ids[:, -CONTEXT_LENGTH:]
-        attention_mask = attention_mask[:, -CONTEXT_LENGTH:]
-
     generate_kwargs = dict(
         input_ids=input_ids,
         attention_mask=attention_mask,
@@ -45,32 +33,30 @@ def predict(message, history, system_prompt, temperature, max_new_tokens, top_k,
         repetition_penalty=repetition_penalty,
         top_p=top_p
     )
+    
     t = Thread(target=model.generate, kwargs=generate_kwargs)
     t.start()
+    
     outputs = []
     for new_token in streamer:
-        outputs.append(new_token)
         if new_token in stop_tokens:
-            
-            break
-        yield "".join(outputs)
-
+            break  # Stop generation but don't add the stop token
+        outputs.append(new_token)
+        yield "".join(outputs).replace("<|im_end|>", "")  # Ensure no leftover stop tokens
 
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_ID)
+model = AutoModelForCausalLM.from_pretrained(MODEL_ID)
 
-# Create Gradio interface
 gr.ChatInterface(
     predict,
-    
-    additional_inputs_accordion=gr.Accordion(label="Parameters", open=False),
     additional_inputs=[
-        gr.Textbox("You are a useful assistant. first recognize user request and then reply carfuly and thinking", label="System prompt"),
+        gr.Textbox("You are a helpful assistant. Format responses clearly using natural Markdown formatting where appropriate.", 
+                 label="System prompt"),
         gr.Slider(0, 1, 0.6, label="Temperature"),
         gr.Slider(0, 4096, 512, label="Max new tokens"),
         gr.Slider(1, 80, 40, label="Top K sampling"),
         gr.Slider(0, 2, 1.1, label="Repetition penalty"),
         gr.Slider(0, 1, 0.95, label="Top P sampling"),
     ],
-).queue().launch()
+    css=".message { white-space: pre-wrap; }",  # Preserve newlines
+).queue().launch()