from transformers import AutoModelForCausalLM, AutoTokenizer
model_id = "ynakazat11/llm-jp-4-8b-instruct-sansu"
tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True, torch_dtype="bfloat16", device_map="auto")
system = (
"あなたは中学受験算数を教える先生です。"
"問題文を読み、小学生にもわかるように、算数の手法(つるかめ算・差集め算・"
"面積図・線分図・比など)を使って段階的に解説してください。"
"文字式・方程式・代数(x, y などの未知数を立てる方法)は使わないでください。"
)
user = "270mの道の端から端まで桜の木を植えます。木と木の間隔を9mにすると、木は何本植えられますか。"
messages = [{"role": "system", "content": system}, {"role": "user", "content": user}]
prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
out = model.generate(**inputs, max_new_tokens=512, temperature=0.3)
print(tokenizer.decode(out[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True))