Update README.md

README.md

@@ -23,19 +23,21 @@ Yamase-12B-v0.1は、[Mistral-Nemo-Instruct](https://huggingface.co/mistralai/Mi
 ```python
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
-B_INST, E_INST = "[INST]", "[/INST]"
+
 text = "旅行に行くと高層ビルがたくさん建っていました。これからどのようなことが推測できますか？"
 model_name = "sudy-super/Yamase-12B"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16)
+model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16,)
 if torch.cuda.is_available():
     model = model.to("cuda")
 model.eval()
-prompt = "{bos_token}{b_inst}{prompt}{e_inst}".format(
-    bos_token=tokenizer.bos_token,
-    b_inst=B_INST,
-    prompt=text,
-    e_inst=E_INST,
+messages = [
+    {"role": "user", "content": text},
+]
+prompt = tokenizer.apply_chat_template(
+    messages,
+    tokenize=False,
+    add_generation_prompt=True,
 )
 with torch.no_grad():
     token_ids = tokenizer.encode(prompt, add_special_tokens=False, return_tensors="pt")