Update README.md

README.md

@@ -16,10 +16,18 @@ Stockmark-100b-instruct-v0.1 is an instruction tuned version of [stockmark-100b]
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
+prompt_template = """### 指示:
+{instruction}
+
+### 応答:
+"""
+
 tokenizer = AutoTokenizer.from_pretrained("stockmark/stockmark-100b")
 model = AutoModelForCausalLM.from_pretrained("stockmark/stockmark-100b", device_map="auto", torch_dtype=torch.bfloat16)
 
-input_ids = tokenizer("生成AIとは？", return_tensors="pt").input_ids.to(model.device)
+instruction = "生成AIとは？"
+prompt = prompt_template.format(instruction=instruction)
+input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(model.device)
 with torch.inference_mode():
     tokens = model.generate(
         input_ids,