Update README.md

README.md

@@ -27,15 +27,19 @@ import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
 tokenizer = AutoTokenizer.from_pretrained("stockmark/Stockmark-2-100B-Instruct-beta")
-model = AutoModelForCausalLM.from_pretrained("stockmark/Stockmark-2-100B-Instruct-beta", device_map="auto", torch_dtype=torch.bfloat16)
+model = AutoModelForCausalLM.from_pretrained(
+    "stockmark/Stockmark-2-100B-Instruct-beta", device_map="auto", torch_dtype=torch.bfloat16
+)
 
-instruction = ""
-input_ids = tokenizer(instruction, add_generation_prompt=True).input_ids.to(model.device)
+instruction = "自然言語処理とは？"
+input_ids = tokenizer.apply_chat_template(
+    [{"role": "user", "content": instruction}], add_generation_prompt=True, return_tensors="pt"
+).to(model.device)
 
 with torch.inference_mode():
     tokens = model.generate(
         input_ids,
-        max_new_tokens = 1024,
+        max_new_tokens = 256,
         do_sample = True,
         temperature = 0.7,
         top_p = 0.95,