File size: 454 Bytes
416f658 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 | from transformers import AutoModelForCausalLM, AutoTokenizer
checkpoint = "HuggingFaceTB/SmolLM2-135M"
# "cpu" for CPU usage
# "gpu" if Torch is compiled with CUDA enabled
device = "cpu"
tokenizer = AutoTokenizer.from_pretrained(checkpoint)
model = AutoModelForCausalLM.from_pretrained(checkpoint).to(device)
inputs = tokenizer.encode("Gravity is", return_tensors="pt").to(device)
outputs = model.generate(inputs)
print(tokenizer.decode(outputs[0]))
|