Change prompt to Llama format

app.py

@@ -18,8 +18,8 @@ tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.1-8B", use_fast=Tr
 vocab_size = len(tokenizer)
 eos_token_id = tokenizer.eos_token_id
 mask_token_id = tokenizer.encode('MASK', add_special_tokens=False)[0]
-# assistant_marker_ids = tokenizer.encode("<|start_header_id|>assistant<|end_header_id|>", add_special_tokens=False)
-assistant_marker_ids = tokenizer.encode("Assistant:", add_special_tokens=False)
+assistant_marker_ids = tokenizer.encode("<|start_header_id|>assistant<|end_header_id|>", add_special_tokens=False)
+# assistant_marker_ids = tokenizer.encode("Assistant:", add_special_tokens=False)
 # def load_model():
 #     ckpt_path = hf_hub_download(
 #         repo_id="ruurd/tini_bi_m",
@@ -195,19 +195,19 @@ def generate_diffusion_text(input_ids, top_p, top_k):
         conf = probs[range(len(sampled)), sampled].cpu().numpy()
     return sampled, conf 
 
-# def format_chat_prompt(question):
-#     return (
-#         "<|begin_of_text|>\n"
-#         "<|start_header_id|>system<|end_header_id|>\n"
-#         "You are a helpful assistant.\n"
-#         "<|start_header_id|>user<|end_header_id|>\n"
-#         f"{question}\n"
-#         "<|start_header_id|>assistant<|end_header_id|>\n"
-#     )
 def format_chat_prompt(question):
-    return(
-        f"User:{question}\nAssistant:"
+    return (
+        "<|begin_of_text|>\n"
+        "<|start_header_id|>system<|end_header_id|>\n"
+        "You are a helpful assistant.\n"
+        "<|start_header_id|>user<|end_header_id|>\n"
+        f"{question}\n"
+        "<|start_header_id|>assistant<|end_header_id|>\n"
     )
+# def format_chat_prompt(question):
+#     return(
+#         f"User:{question}\nAssistant:"
+#     )
 
 
 # --- Inference Wrapper ---