Update README.md

README.md

@@ -22,11 +22,15 @@ To load and use this model with Hugging Face Transformers:
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
-model_name = "maomaocun/LLaDA-Prometheus-no-template"
+model_name = "maomaocun/LLaDA-Prometheus"
 tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
 model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16, trust_remote_code=True).to("cuda")
 
-prompt = "Can you tell me an engaging short story about a brave young astronaut who discovers an ancient alien civilization on a distant planet? Make it adventurous and heartwarming, with a twist at the end."
+# 使用对话模板
+messages = [
+    {"role": "user", "content": "Can you tell me an engaging short story about a brave young astronaut who discovers an ancient alien civilization on a distant planet? Make it adventurous and heartwarming, with a twist at the end."}
+]
+prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
 
 inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
 input_ids = inputs['input_ids']
@@ -38,10 +42,10 @@ for chunk in model.generate(
     block_length=64,
     threshold=0.9,
     streaming=True,
-    eos_token_id=tokenizer.eos_token,
+    eos_token_id=126348  # 注意：这里改为 eos_token_id
 ):
     all_generated_ids = torch.cat([input_ids, chunk], dim=-1)
-    text = tokenizer.batch_decode(all_generated_ids, skip_special_tokens=False)[0].split(tokenizer.eos_token)[0]
+    text = tokenizer.batch_decode(all_generated_ids, skip_special_tokens=True)[0]  # 改为 skip_special_tokens=True，并移除手动 split
     print(text, end='', flush=True)
 ```