Update app.py

app.py

@@ -16,6 +16,11 @@ tokenizer = AutoTokenizer.from_pretrained(model_name)
 if tokenizer.pad_token is None:
     tokenizer.pad_token = tokenizer.eos_token
 
+# Define a custom chat template if one is not available
+if tokenizer.chat_template is None:
+    # Basic ChatML-style template
+    tokenizer.chat_template = "{% for message in messages %}\n{% if message['role'] == 'system' %}<|system|>\n{{ message['content'] }}\n{% elif message['role'] == 'user' %}<|user|>\n{{ message['content'] }}\n{% elif message['role'] == 'assistant' %}<|assistant|>\n{{ message['content'] }}\n{% endif %}\n{% endfor %}\n{% if add_generation_prompt %}<|assistant|>\n{% endif %}"
+
 def respond(
     message,
     history: list[tuple[str, str]],
@@ -34,11 +39,11 @@ def respond(
             messages.append({"role": "assistant", "content": assistant_msg})
     
     messages.append({"role": "user", "content": message})
-
+    
     # Format prompt using chat template
     prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
     inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-
+    
     # Set up streaming
     streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
     
@@ -53,11 +58,11 @@ def respond(
         do_sample=do_sample,
         pad_token_id=tokenizer.pad_token_id
     )
-
+    
     # Start generation in separate thread
     thread = Thread(target=model.generate, kwargs=generation_kwargs)
     thread.start()
-
+    
     # Stream response
     response = ""
     for token in streamer: