Update app.py

app.py

@@ -16,10 +16,14 @@ def patched_get_type(schema):
 client_utils.get_type = patched_get_type
 
 # Load model and tokenizer
-model_name = "arshiaafshani/Arsh-llm" 
+model_name = "arshiaafshani/Arsh-llm"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32)
 
+# تنظیم توکن‌های خاص
+tokenizer.bos_token = "<sos>"
+tokenizer.eos_token = "<|endoftext|>"
+
 # Create pipeline
 pipe = pipeline(
     "text-generation",
@@ -29,8 +33,12 @@ pipe = pipeline(
 )
 
 def respond(message, chat_history, system_message, max_tokens, temperature, top_p, top_k, repeat_penalty):
-    # Prepare prompt
-    prompt = f"{system_message}\n\nUser: {message}\nAssistant:"
+    # Prepare prompt using apply_chat_template
+    chat_history = chat_history or []
+    messages = [{"role": "system", "content": system_message}] + \
+               [{"role": "user", "content": msg} for msg, _ in chat_history] + \
+               [{"role": "user", "content": message}, {"role": "assistant", "content": ""}]
+    prompt = tokenizer.apply_chat_template(messages, tokenize=False)
     
     # Generate response
     output = pipe(
@@ -44,7 +52,8 @@ def respond(message, chat_history, system_message, max_tokens, temperature, top_
         pad_token_id=tokenizer.eos_token_id
     )
     
-    response = output[0]['generated_text'].split("Assistant:")[-1].strip()
+    # Extract response
+    response = output[0]['generated_text'][len(prompt):].strip()
     
     # Update chat history
     chat_history.append((message, response))
@@ -80,4 +89,4 @@ with gr.Blocks() as demo:
     clear.click(lambda: None, None, chatbot, queue=False)
 
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860, share=True)
+    demo.launch(server_name="0.0.0.0", server_port=7860)