from transformers import AutoModelForCausalLM, AutoTokenizer
repo_id = "YuenWu/Qwen3-0.6B-wyc"
tokenizer = AutoTokenizer.from_pretrained(repo_id)
model = AutoModelForCausalLM.from_pretrained(
repo_id,
torch_dtype="auto",
device_map="auto",
)
messages = [
{"role": "user", "content": "Solve: 2 + 3."},
]
text = tokenizer.apply_chat_template(
messages,
tokenize=False,
add_generation_prompt=True,
enable_thinking=True,
)
inputs = tokenizer([text], return_tensors="pt").to(model.device)
generated_ids = model.generate(**inputs, max_new_tokens=256)
new_tokens = generated_ids[:, inputs.input_ids.shape[1]:]
response = tokenizer.batch_decode(new_tokens, skip_special_tokens=True)[0]
print(response)