import torchfrom transformers import AutoModelForCausalLM, AutoTokenizer model_id = "YOUR_USERNAME/SmolLM2-135M-Reasoning-5K" tokenizer = AutoTokenizer.from_pretrained(model_id)model = AutoModelForCausalLM.from_pretrained( model_id, torch_dtype="auto", device_map="auto",) messages = [ { "role": "system", "content": 'You are a helpful AI assistant. For difficult problems, reason carefully inside <think> and </think> tags, then provide a clear final answer.', }, { "role": "user", "content": "A farmer has 17 sheep. All but 9 run away. How many remain?", },] inputs = tokenizer.apply_chat_template( messages, tokenize=True, add_generation_prompt=True, return_tensors="pt", return_dict=True,).to(model.device) with torch.inference_mode(): output = model.generate( **inputs, max_new_tokens=256, do_sample=False, repetition_penalty=1.05, ) new_tokens = output[0, inputs["input_ids"].shape[1]:]print(tokenizer.decode(new_tokens, skip_special_tokens=False))