from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

ckpt = "facebook/MobileLLM-R1-950M"
tok = AutoTokenizer.from_pretrained(ckpt)
model = AutoModelForCausalLM.from_pretrained(
 ckpt,
 torch_dtype=torch.float16, # or int8/int4 via bitsandbytes/AutoGPTQ
 device_map="auto"
)

prompt = "Solve 48/6 + 7*3. Show steps briefly."
inputs = tok(prompt, return_tensors="pt").to(model.device)

with torch.inference_mode:
 out = model.generate(
 **inputs,
 max_new_tokens=160,
 temperature=0.2,
 do_sample=False
 )

print(tok.decode(out[0], skip_special_tokens=True))